{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9d3ff43",
   "metadata": {},
   "outputs": [],
   "source": [
    "import scanpy as sc\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import pickle\n",
    "import pandas as pd\n",
    "%matplotlib inline\n",
    "import seaborn as sns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8dde3885",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Utility functions and classes for cross species\n",
    "analysis\n",
    "\n",
    "@yhr91\n",
    "\"\"\"\n",
    "\n",
    "from sklearn.metrics import euclidean_distances\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib\n",
    "from collections import defaultdict\n",
    "from sklearn.metrics.pairwise import cosine_distances\n",
    "from scipy.stats import spearmanr\n",
    "import plotly.express as px\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import warnings\n",
    "import scanpy as sc\n",
    "from sklearn.metrics import adjusted_mutual_info_score, rand_score\n",
    "from collections import Counter\n",
    "from scipy.stats import mode\n",
    "import operator\n",
    "\n",
    "# --------\n",
    "\n",
    "class cross_species_acc():\n",
    "    \"\"\"\n",
    "    Class for calculating cross species accuracy metrics\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, adata, base_species='human', \n",
    "                 target_species='mouse', label_col='CL_class_coarse', \n",
    "                 metric='cosine', medoid=False, space='raw'):\n",
    "\n",
    "        self.adata = adata\n",
    "        self.base_species = base_species\n",
    "        self.target_species = target_species\n",
    "        self.label_col = label_col\n",
    "        self.metric = metric\n",
    "        self.medoid=medoid\n",
    "        self.space=space\n",
    "        \n",
    "        # Calculate accuracy metrics\n",
    "        self.calc_cross_species_label_matches()\n",
    "\n",
    "    def find_all_species_centres(self):\n",
    "        \"\"\"\n",
    "        Finds all species-specific centroids given an AnnData object\n",
    "        \"\"\"\n",
    "        if self.space == 'umap':\n",
    "            key = 'X_umap'\n",
    "        elif self.space == 'samap':\n",
    "            key = 'X_umap_samap'\n",
    "        elif self.space == 'scanorama':\n",
    "            key = 'X_scanorama'\n",
    "        elif self.space == 'harmony':\n",
    "            key = 'X_harmony'\n",
    "        \n",
    "        centres = {}\n",
    "        centres['size'] = {}\n",
    "        base_cluster_sizes = {}\n",
    "        for species in self.adata.obs['species'].unique():\n",
    "            centres[species] = {}\n",
    "            species_set = self.adata[self.adata.obs['species']==species]\n",
    "            \n",
    "            for l in species_set.obs[self.label_col].unique():\n",
    "                subset = species_set[species_set.obs[self.label_col] == l]\n",
    "                \n",
    "                # If space is not raw then use the right obsm column\n",
    "                if self.space != 'raw':\n",
    "                    subset_data = subset.obsm[key]\n",
    "                else:\n",
    "                    subset_data = subset.X.toarray()\n",
    "\n",
    "                # Deal with exceptions\n",
    "                if len(subset)<1:\n",
    "                    continue\n",
    "                elif len(subset)==1:\n",
    "                    centres[species][l] = subset_data[0]\n",
    "                \n",
    "                # Use centroid or medoid\n",
    "                centroid = np.mean(subset_data, 0)\n",
    "                if self.medoid:\n",
    "                    centres[species][l] =\\\n",
    "                        self.get_medoid(subset_data, centroid)\n",
    "                else:\n",
    "                    centres[species][l] = centroid\n",
    "                    \n",
    "                # This is for normalization of distances\n",
    "                if species == self.base_species:\n",
    "                    dist_mat = euclidean_distances(subset_data)\n",
    "                    centres['size'][l] = np.max(dist_mat)\n",
    "            \n",
    "        return pd.DataFrame(centres).dropna()\n",
    "\n",
    "    \n",
    "    def calc_cross_species_label_matches(self):\n",
    "        \"\"\"\n",
    "        Given Anndata object, returns:\n",
    "        - matches: number of cluster centres in base species that \n",
    "        have the same cluster label in the target species as nn\n",
    "        - dist: 'normalized' distance between cluster centre of base \n",
    "        species and the cluster centre with the same label in target \n",
    "        species\n",
    "        \n",
    "        TODO: This is not generalized to more than 2 species\n",
    "        \"\"\"\n",
    "        warnings.filterwarnings(\"ignore\")\n",
    "        \n",
    "        centres = self.find_all_species_centres()\n",
    "        dist = 0\n",
    "        norm_dist = 0\n",
    "        matches = 0\n",
    "        matches_names = []\n",
    "        matches_names_all = []\n",
    "        target_centres = np.vstack(\n",
    "                   centres.loc[:,self.target_species].values)\n",
    "        \n",
    "        for idx, ctype in enumerate(centres.index):\n",
    "            base = centres.loc[ctype, self.base_species]\n",
    "            base_targets = np.vstack([base, target_centres])\n",
    "\n",
    "            if self.metric=='cosine':\n",
    "                distances = cosine_distances(base_targets)[0][1:]\n",
    "                \n",
    "            pred_match = np.argmin(distances)\n",
    "            if  pred_match == idx:\n",
    "                matches += 1\n",
    "                matches_names.append(ctype)\n",
    "            \n",
    "            matches_names_all.append((ctype,\n",
    "                         centres.index[pred_match]))\n",
    "            dist += distances[idx]\n",
    "            norm_dist += distances[idx]/centres.loc[ctype, 'size']\n",
    "\n",
    "        self.cross_species_label_dist = dist\n",
    "        self.cross_species_label_norm_dist = norm_dist\n",
    "        self.cross_species_label_matches = matches\n",
    "        self.cross_species_label_matches_names = matches_names\n",
    "        self.cross_species_label_matches_names_all = matches_names_all\n",
    "        \n",
    "        warnings.filterwarnings(\"always\")\n",
    "     \n",
    "    \n",
    "    def get_medoid(self, data, centroid):\n",
    "        dists = euclidean_distances(np.vstack([centroid,data]))[0]\n",
    "        return data[np.argsort(dists)[1]-1]\n",
    "        \n",
    "        \n",
    "# --------\n",
    "\n",
    "\n",
    "class embedding_CL_comparison():\n",
    "    \"\"\"\n",
    "    Class for comparing embedding with cell ontology\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, adata, label_col='CL_class_coarse', CL_ID_col='CL_ID_coarse',\n",
    "                 metric='cosine', features='raw'):\n",
    "\n",
    "        warnings.filterwarnings(\"ignore\")\n",
    "        self.adata = adata\n",
    "        self.label_col = label_col\n",
    "        self.CL_ID_col = CL_ID_col\n",
    "        self.metric = metric\n",
    "        self.features = features\n",
    "        self.labels = self.adata.obs[self.label_col].unique()\n",
    "        self.centres = []\n",
    "        self.centres_ranked = []\n",
    "        self.CL_centres_ranked = []\n",
    "\n",
    "        # Get centres, nns and ranks\n",
    "        self.get_centre_ranks()\n",
    "        self.get_CL_ranks()\n",
    "\n",
    "        # Calculate metrics\n",
    "        self.spearman_corr = {}\n",
    "        self.hits_at_k = {}\n",
    "        for id_ in self.labels:         \n",
    "            self.spearman_corr[id_] = spearmanr(self.CL_centres_ranked[id_], \n",
    "                self.centres_ranked[id_])[0]\n",
    "            self.hits_at_k[id_] = self.get_hits_topk(self.CL_centres_ranked[id_], \n",
    "                self.centres_ranked[id_])\n",
    "            \n",
    "        warnings.filterwarnings(\"always\")\n",
    "\n",
    "        \n",
    "    # Implement cluster centroid\n",
    "\n",
    "    def find_centre(self, cluster, medioid=False):\n",
    "        \"\"\"\n",
    "        Find cluster centre: either centroid or medioid\n",
    "        \"\"\"\n",
    "        if medioid:\n",
    "            dist = euclidean_distances(cluster)\n",
    "            medioid = np.argmin(dist.sum(0))\n",
    "            return cluster[medioid].toarray()\n",
    "\n",
    "        else:\n",
    "            return np.mean(cluster,0)\n",
    "    \n",
    "    \n",
    "    def get_outlier_idx(self, CL_centres, centres, k=10):\n",
    "        \"\"\"\n",
    "        Get top or bottom ranked nn\n",
    "        \"\"\"\n",
    "        outliers = []\n",
    "        for i, pair in enumerate(list(zip(CL_centres, centres))):\n",
    "            if pair[0] < k or pair[1] < k:\n",
    "                outliers.append(i)\n",
    "        return outliers\n",
    "\n",
    "    \n",
    "    def get_hits_topk(self, CL_centres, centres, k=10):\n",
    "        \"\"\"\n",
    "        Get numbers of matches within top k\n",
    "        \"\"\"\n",
    "        return len(set(CL_centres[:k]).intersection(set(centres[:k])))\n",
    "        \n",
    "\n",
    "    def get_centre_ranks(self):\n",
    "        \"\"\"\n",
    "        Get nn ranks for cluster centres\n",
    "        \"\"\"\n",
    "        for cell_type in self.labels:\n",
    "            if self.features=='raw':\n",
    "                self.centres.append(self.find_centre(\n",
    "                    self.adata[self.adata.obs[self.label_col] == cell_type].X))\n",
    "        self.centres = np.vstack(self.centres)\n",
    "                \n",
    "        if self.metric=='euclidean':\n",
    "            centres_dist = euclidean_distances(self.centres)\n",
    "\n",
    "        if self.metric=='cosine':\n",
    "            centres_dist = cosine_distances(self.centres)\n",
    "\n",
    "        self.centres_ranked = {k:v for k,v in zip(\n",
    "            self.labels, np.argsort(centres_dist))}\n",
    "        \n",
    "    def get_CL_ranks(self):\n",
    "        \"\"\"\n",
    "        Get nn ranks for cell ontology cluster centres\n",
    "        \"\"\"\n",
    "        all_CL_distances = pd.read_csv('/dfs/project/cross-species/data/lung/shared/CL_similarity_RW.csv',\n",
    "                                       index_col=0)\n",
    "\n",
    "        CL_sim_matrix = all_CL_distances.loc[self.adata.obs[self.CL_ID_col].unique(),\n",
    "                                             self.adata.obs[self.CL_ID_col].unique()]\n",
    "        \n",
    "        ID_dict = self.adata.obs.set_index(self.label_col).to_dict()['CL_ID_coarse']\n",
    "        inv_ID_dict = {v: k for k, v in ID_dict.items()}\n",
    "\n",
    "        self.CL_centres_ranked = {inv_ID_dict[k]:v for k,v in zip(self.adata.obs['CL_ID_coarse'].unique(),\n",
    "                                                           np.argsort(-CL_sim_matrix.values))}\n",
    "\n",
    "    \n",
    "    def plot_rank_scatter(self):\n",
    "        \"\"\"\n",
    "        Create rank scatter plot between embedding nn and CL nn\n",
    "        \"\"\"\n",
    "        fig, axs = plt.subplots(5, 5, sharex=True, sharey=True, figsize=[20,15])\n",
    "        it = 0\n",
    "        spearman_corr = {}\n",
    "        hits_at_k = {}\n",
    "        outlier=False\n",
    "\n",
    "        for i in range(5):\n",
    "            for j in range(5):\n",
    "                if it == len(self.labels):\n",
    "                    break\n",
    "\n",
    "                id_ = self.labels[it]\n",
    "                if outlier:\n",
    "                    outlier_idx = get_outlier_idx(self.CL_centres_ranked[id_], \n",
    "                        self.centres_ranked[id_])\n",
    "                    axs[i, j].scatter(self.CL_centres_ranked[id_][outlier_idx], \n",
    "                    self.centres_ranked[id_][outlier_idx])\n",
    "                else:\n",
    "                    axs[i, j].scatter(self.CL_centres_ranked[id_], \n",
    "                    self.centres_ranked[id_])\n",
    "                axs[i, j].set_title(id_)\n",
    "                axs[i, j].plot([0,32],[0,32], 'k')   \n",
    "                it += 1\n",
    "                \n",
    "    def plot_hits_at_k(self):\n",
    "        plot_df = pd.DataFrame.from_dict(self.hits_at_k, orient='index')\n",
    "        plot_df = plot_df.rename(columns={0:'Value'})\n",
    "        plot_df = plot_df.sort_values('Value')\n",
    "\n",
    "        plt.figure(figsize=[8,12])\n",
    "        plt.barh(plot_df.index, plot_df['Value'])\n",
    "        plt.ylabel('Cell Type')\n",
    "        plt.xlabel('Hits @ k')\n",
    "        plt.title('Hits @ k (Embedding space compared to Cell Ontology)')\n",
    "        plt.xlim([0,10])\n",
    "        \n",
    "        \n",
    "    def plot_spearman(self):\n",
    "        plot_df = pd.DataFrame.from_dict(self.spearman_corr, orient='index')\n",
    "        plot_df = plot_df.rename(columns={0:'Value'})\n",
    "        plot_df = plot_df.sort_values('Value')\n",
    "\n",
    "        plt.figure(figsize=[8,12])\n",
    "        plt.barh(plot_df.index, plot_df['Value'])\n",
    "        plt.ylabel('Cell Type')\n",
    "        plt.xlabel('Spearman Correlation')\n",
    "        plt.title('Spearman Correlation (Embedding space compared to Cell Ontology)')\n",
    "        plt.xlim([-1,1])\n",
    "          \n",
    "        \n",
    "# --------\n",
    "\n",
    "## KNN analysis per cell\n",
    "## TODO integrate these functions into cross_species_acc class\n",
    "\n",
    "def get_knn_label(cell_names, adata, col):\n",
    "    \"\"\"\n",
    "    Returns majority class labels of nearest neighbors. \n",
    "    Will return random label in case of tie\n",
    "    \"\"\"\n",
    "    \n",
    "    return adata[cell_names].obs[col].value_counts().index[0]\n",
    "\n",
    "\n",
    "def cross_species_knn_all(adata, k=1, species='human', space='raw',\n",
    "                          col = 'cell_type', metric='euclidean',\n",
    "                          verbose = False, consider_same_species=False):\n",
    "    \"\"\"Runs cross species k nearest neighbor on all cells\n",
    "    \"\"\"\n",
    "\n",
    "    # Create distance matrix\n",
    "    if space == 'raw':\n",
    "        X = adata.X\n",
    "    elif space == 'umap':\n",
    "        X = adata.obsm['X_umap']\n",
    "    elif space == 'samap':\n",
    "        X = adata.obsm['X_umap_samap']\n",
    "    elif space == 'scanorama':\n",
    "        X = adata.obsm['X_scanorama']\n",
    "    elif space == 'harmony':\n",
    "        X = adata.obsm['X_harmony']\n",
    "    \n",
    "    # Slow step\n",
    "    if metric == 'euclidean':\n",
    "        dist_mat = euclidean_distances(X)\n",
    "    elif metric == 'cosine':\n",
    "        dist_mat = cosine_distances(X)\n",
    "        \n",
    "    if consider_same_species:\n",
    "        # Get indices for species and nonspecies cells\n",
    "        species_idx = np.where(adata.obs['species']==species)[0]\n",
    "        adata.obs['temp_label'] = adata.obs['species'].astype(str) +\\\n",
    "                            '_' + adata.obs[col].astype(str)\n",
    "\n",
    "        nns = []\n",
    "        for idx in species_idx:\n",
    "            curr_temp_label = adata.obs['temp_label'][idx]\n",
    "            row = dist_mat[idx,:]\n",
    "            possible_nbrs = np.where(adata.obs['temp_label'] != curr_temp_label)[0]\n",
    "            \n",
    "            row = row[possible_nbrs]\n",
    "            nns.append(possible_nbrs[np.argpartition(row, k)[:k]])\n",
    "            \n",
    "        nbrs = [(adata.obs[col][x], Mode(adata.obs[col][y].astype('str').values)) \n",
    "                        for x,y in zip(species_idx, nns)]\n",
    "    \n",
    "    else:\n",
    "        # Get indices for species and nonspecies cells\n",
    "        species_idx = np.where(adata.obs['species']==species)[0]\n",
    "        nonspecies_idx = np.where(adata.obs['species'] != species)[0]\n",
    "\n",
    "        # Slow step\n",
    "        reduced_dist_mat = dist_mat[species_idx,:][:,nonspecies_idx]\n",
    "        \n",
    "        nns = [list(nonspecies_idx[y]) \n",
    "               for y in np.argpartition(reduced_dist_mat, k)[:,:k]]\n",
    "\n",
    "        nbrs = [(adata.obs[col][x], Mode(adata.obs[col][y].astype('str').values)) \n",
    "                        for x,y in zip(species_idx, nns)]\n",
    "    \n",
    "    return nbrs\n",
    "\n",
    "\n",
    "def cluster_knn(cluster_knn_df, label):\n",
    "    \"\"\"\n",
    "    Given majority k nearest cross species neigbhor class for each cell, \n",
    "    identifies the k nearest neighbors for the given cluster\n",
    "    \"\"\"\n",
    "    \n",
    "    list_ = list(cluster_knn_df[cluster_knn_df['Source_Cell']==label].\n",
    "                value_counts().items())\n",
    "    x = pd.DataFrame([(s[0][1],s[1]) for s in list_])\n",
    "    x['Source_Cluster'] = list_[0][0][0]\n",
    "    x = x.rename(columns={0:'Cross_Species_KNN_Label', 1:'Score'})\n",
    "    x['Score'] = x['Score']/x['Score'].sum()\n",
    "    return x\n",
    "\n",
    "\n",
    "def cluster_knn_all(all_nbrs):\n",
    "    \"\"\"\n",
    "    Given majority k nearest cross species neigbhor class for each cell, \n",
    "    identifies the k nearest neighbors for all clusters\n",
    "    \"\"\"\n",
    "    cluster_knn_df = pd.DataFrame(all_nbrs)\n",
    "    cluster_knn_df = cluster_knn_df.rename(columns={0:'Source_Cell',1:'Cross_Species_KNN'})\n",
    "\n",
    "    return [cluster_knn(cluster_knn_df, c) \n",
    "         for c in cluster_knn_df['Source_Cell'].unique()]   \n",
    "\n",
    "\n",
    "def plot_cluster_knn_bar(df, source='human', other='mouse', ax=None, title=None):\n",
    "    \"\"\"\n",
    "    Creates a stacked bar plot to identify majority k nearest neighbors for\n",
    "    a given cluster, on a specific axis\n",
    "    \"\"\"\n",
    "    if ax is None:\n",
    "        fig, (ax) = plt.subplots(1, 1, sharex=True, sharey=True, figsize=(6, 17), frameon=False)\n",
    "    bars = defaultdict(int)\n",
    "    colors = defaultdict(int)\n",
    "    tick = -1; \n",
    "    tick_pos = {}\n",
    "    colors_list = ['','r','g','b','y','k','purple','g','b','y','k','g','b','y',\n",
    "                                                       'k','g','b','y','k',\n",
    "                                                       'k','g','b','y','k',\n",
    "                                                       'k','g','b','y','k'] + [\"w\"]*50\n",
    "    df = df.sort_values(['Source_Cluster', 'Score'], ascending=False)\n",
    "\n",
    "    for i in df.iterrows():\n",
    "        color = 'k'\n",
    "        x = i[1]['Source_Cluster']\n",
    "        y = i[1]['Score']        \n",
    "\n",
    "        if i[1]['Cross_Species_KNN_Label'] == i[1]['Source_Cluster']:\n",
    "            color= 'w'\n",
    "        left = bars[x]\n",
    "        bars[x] = bars[x] + y\n",
    "        colors[x] = colors[x] + 1\n",
    "        if colors_list[colors[x]] == 'r':    \n",
    "            tick = tick+1\n",
    "            tick_pos[x] = tick\n",
    "            if bars[x]>=0.50:\n",
    "                ax.text(0.1,tick-0.2,i[1]['Cross_Species_KNN_Label'], color=color)\n",
    "        ax.barh(tick_pos[x], y, left=left, color=colors_list[colors[x]], alpha=0.5)\n",
    "        \n",
    "    keys = sorted(tick_pos.keys())\n",
    "    vals = [tick_pos[k] for k in keys]\n",
    "    ax.set_yticks(vals, keys)\n",
    "    ax.set_ylabel(source)\n",
    "    ax.set_xlabel('Percentage of cells with cross-species KNN class')\n",
    "    if title is not None:\n",
    "        ax.set_title(title)\n",
    "    \n",
    "    return \n",
    "\n",
    "## ---------------------------------\n",
    "## Alignment scores\n",
    "## -----------------------------------\n",
    "\n",
    "def alignment_score(fname, col='cell_type', space='raw', k=1,\n",
    "                    species='human', consider_same_species=False):\n",
    "    adata = sc.read_h5ad(fname)\n",
    "    if space=='umap':\n",
    "        sc.pp.pca(adata, n_comps=50)\n",
    "        sc.pp.neighbors(adata, n_neighbors=15)\n",
    "        sc.tl.umap(adata)\n",
    "    all_nbrs = cross_species_knn_all(adata, col=col, metric='cosine', space=space, k=k,\n",
    "                                     species=species, consider_same_species=consider_same_species)\n",
    "    all_cluster_nbrs  = [i for i in cluster_knn_all(all_nbrs)]\n",
    "    all_cluster_nbrs = pd.concat(all_cluster_nbrs).reset_index(drop=True) \n",
    "    return all_cluster_nbrs\n",
    "\n",
    "def compare_matches(alignments, true_alignments):\n",
    "    aligns = alignments.merge(true_alignments, on=['Source_Cluster', 'Cross_Species_KNN_Label'])\n",
    "    #aligns = aligns.merge(true_alignments, on=['Source_Cluster'], how='outer').fillna(0)\n",
    "    \n",
    "    return aligns\n",
    "\n",
    "def score_matches(alignments, true_alignments, thresh=0.5, ret_matches=False):\n",
    "    matches = compare_matches(alignments, true_alignments)\n",
    "    if ret_matches:\n",
    "        return matches[matches['Score']>thresh]\n",
    "    else:\n",
    "        return sum(matches['Score']>thresh)\n",
    "    return\n",
    "\n",
    "def create_comparison_plot_df(all_cluster_nbrs1, all_cluster_nbrs2, true_map):\n",
    "    \n",
    "    knn_scores_1 = compare_matches(all_cluster_nbrs1, true_map)\n",
    "    knn_scores_2 = compare_matches(all_cluster_nbrs2, true_map)\n",
    "\n",
    "    plot_df = knn_scores_1.merge(knn_scores_2, on='Source_Cluster') \n",
    "    return plot_df\n",
    "\n",
    "def get_comparison_plot(plot_df, bars=2, labels_=['Method1', 'Method2']):\n",
    "    plt.figure(figsize=[20,5])\n",
    "    ax=plt.gca()\n",
    "\n",
    "    labels = plot_df['Source_Cluster'].values\n",
    "    x = np.arange(len(labels))  # the label locations\n",
    "    width = 0.35  # the width of the bars\n",
    "\n",
    "    #fig, ax = plt.subplots()\n",
    "    rects1 = ax.bar(x - width/2, plot_df['Score_x'].values, width, label=labels_[0])\n",
    "    if bars == 2:\n",
    "        rects2 = ax.bar(x + width/2, plot_df['Score_y'].values, width, label=labels_[1])\n",
    "\n",
    "    # Add some text for labels, title and custom x-axis tick labels, etc.\n",
    "    ax.set_ylabel('% of nearest neighbors of\\n correct cross species label')\n",
    "    ax.set_title('Cross species cell type alignment')\n",
    "    plt.xticks(x, labels, rotation='vertical')\n",
    "    ax.set_xticklabels(labels)\n",
    "    ax.legend() \n",
    "\n",
    "    \n",
    "def get_cell_alignment(align_df, true_df, count_df):\n",
    "    align_df['Source_Cluster'] = align_df['Source_Cluster'].astype('str')\n",
    "    count_df['Source_Cluster'] = count_df['Source_Cluster'].astype('str')\n",
    "    true_df['Source_Cluster'] = true_df['Source_Cluster'].astype('str')\n",
    "    align_df['Cross_Species_KNN_Label'] = align_df['Cross_Species_KNN_Label'].astype('str')                  \n",
    "    true_df['Cross_Species_KNN_Label'] = true_df['Cross_Species_KNN_Label'].astype('str')\n",
    "    \n",
    "    df = align_df.merge(true_df, on=['Source_Cluster', 'Cross_Species_KNN_Label'])\n",
    "    df = df.merge(count_df, on='Source_Cluster', how='outer').fillna(0)\n",
    "    return sum((df['Score']*df['count'])/sum(df['count']))\n",
    "\n",
    "\n",
    "def get_alignment_metrics(fname, out_label = 'labels2', orig_label='CL_class_coarse',\n",
    "                          space='raw', species=['human','mouse'], k=1,\n",
    "                 true_labels_path=None, ret_matches = False, consider_same_species=False):\n",
    "    \"\"\"\n",
    "    Function for computing evaluation metrics for embedding\n",
    "    \n",
    "    Outputs:\n",
    "    - species1_nn: Number of cross-species label matches (species 1)\n",
    "    - species2_nn: Number of cross-species label matches (species 2)\n",
    "    - union_nn: Number of cross-species label matches in either species \n",
    "    - mutual_nn: Number of cross-species label matches in both species \n",
    "    - cell_score1: Percentage of cells in species 1 with cross species nn of correct label\n",
    "    - cell_score2: Percentage of cells in species 2 with cross species nn of correct label\n",
    "    - cell_score_combine: Percentage of cells in both species with cross species nn of correct label\n",
    "    - centroid_matches_species1: Number of species 1 centroids that are nn with correct species 2 centroid\n",
    "    - centroid_matches_species2: Number of species 2 centroids that are nn with correct species 1 centroid\n",
    "    - centroid_matches_union: Union of centroid lists\n",
    "    - medoid_matches_species1: Number of species 1 medoids that are nn with correct species 2 medoid\n",
    "    - medoid_matches_species2: Number of species 2 medoids that are nn with correct species 1 medoid\n",
    "    - medoid_matches_union: Union of medoid lists\n",
    "    \"\"\"\n",
    "    \n",
    "    # TODO: This is a very ugly function that needs to be made into a class alongwith\n",
    "    # the functions above it\n",
    "        \n",
    "    # Get cross-species only alignments\n",
    "    alignments = []\n",
    "    print('Finding nns for species 1')\n",
    "    alignments.append(alignment_score(fname, out_label, space=space, species=species[0], k=k,\n",
    "                                     consider_same_species=consider_same_species))\n",
    "    print('Finding nns for species 2')\n",
    "    alignments.append(alignment_score(fname, out_label, space=space, species=species[1], k=k,\n",
    "                                     consider_same_species=consider_same_species))\n",
    "    \n",
    "    # Get true labels\n",
    "    if true_labels_path is None:\n",
    "        if orig_label == 'CL_class_coarse':\n",
    "            true_labels_path = '/dfs/project/cross-species/data/lung/shared/true_CL_class_coarse.csv'\n",
    "        elif orig_label == 'cell_type':\n",
    "            true_labels_path = '/dfs/project/cross-species/data/lung/shared/true_cell_type.csv'\n",
    "        else:\n",
    "            print(\"ERROR: True labels unavailable for this column!, Please set manually\")\n",
    "            return\n",
    "    \n",
    "    true_labels = pd.read_csv(true_labels_path, index_col=0)\n",
    "    cols = []\n",
    "    results = {}\n",
    "    cols.append([c for c in true_labels.columns if species[0] in c][0])\n",
    "    cols.append([c for c in true_labels.columns if species[1] in c][0])\n",
    "\n",
    "    true_dfs = []\n",
    "    true_dfs.append(true_labels.rename(columns={\n",
    "        cols[0]:'Source_Cluster', cols[1]:'Cross_Species_KNN_Label'}))\n",
    "    true_dfs.append(true_labels.rename(columns={\n",
    "        cols[1]:'Source_Cluster', cols[0]:'Cross_Species_KNN_Label'}))\n",
    "\n",
    "    # Score matches\n",
    "    matches = []\n",
    "    matches.append(score_matches(alignments[0], \n",
    "                    true_dfs[0], thresh=0.5, ret_matches=True))\n",
    "    matches.append(score_matches(alignments[1], \n",
    "                    true_dfs[1], thresh=0.5, ret_matches=True))\n",
    "        \n",
    "    for m in matches:\n",
    "        m = m.rename(columns = {'Cross_Species_KNN_Label_x':'Cross_Species_KNN_Label'})\n",
    "        m = m.loc[:,['Score', 'Source_Cluster', 'Cross_Species_KNN_Label']]\n",
    "    results['species1_nn'] = len(matches[0])\n",
    "    results['species2_nn'] = len(matches[1])\n",
    "\n",
    "    # Combine matches\n",
    "    all_matches = matches[0].merge(matches[1], \n",
    "                         left_on=['Source_Cluster', 'Cross_Species_KNN_Label'],\n",
    "                         right_on=['Cross_Species_KNN_Label', 'Source_Cluster'], how='outer')\n",
    "    results['union_nn'] = len(all_matches)\n",
    "    results['mutual_nn'] = len(matches[0].merge(matches[1], \n",
    "                         left_on=['Source_Cluster', 'Cross_Species_KNN_Label'],\n",
    "                         right_on=['Cross_Species_KNN_Label', 'Source_Cluster'], how='inner'))\n",
    "\n",
    "    # Get per-cell alignment scores\n",
    "    adata = sc.read_h5ad(fname)\n",
    "    adata = adata[adata.obs['species'].isin(species)]\n",
    "    ratio1 = sum(adata.obs['species']==species[0])/len(adata)\n",
    "    ratio2 = sum(adata.obs['species']==species[1])/len(adata)\n",
    "\n",
    "    adata1 = adata[adata.obs['species']==species[0]]\n",
    "    adata2 = adata[adata.obs['species']==species[1]]\n",
    "    count_dfs = []\n",
    "    count_dfs.append(pd.DataFrame(adata1.obs[out_label].value_counts()).reset_index().rename(\n",
    "                    columns={'index':'Source_Cluster', out_label:'count'}))\n",
    "    count_dfs.append(pd.DataFrame(adata2.obs[out_label].value_counts()).reset_index().rename(\n",
    "                    columns={'index':'Source_Cluster', out_label:'count'}))\n",
    "\n",
    "    cell_scores = []\n",
    "    cell_scores.append(get_cell_alignment(alignments[0], true_dfs[0], count_dfs[0]))\n",
    "    cell_scores.append(get_cell_alignment(alignments[1], true_dfs[1], count_dfs[1]))\n",
    "    results['cell_score1'] = cell_scores[0]\n",
    "    results['cell_score2'] = cell_scores[1]\n",
    "    results['cell_score_combine'] = ratio1*cell_scores[0] + ratio2*cell_scores[1]\n",
    "\n",
    "    # Get centroid nn score:\n",
    "    for centre,flag in [('centroid', False), ('medoid', True)]:\n",
    "        c_nn1 = cross_species_acc(adata, base_species=species[0], target_species=species[1], \n",
    "                          label_col=out_label, medoid=flag, space=space)\n",
    "        c_nn2 = cross_species_acc(adata, base_species=species[1], target_species=species[0], \n",
    "                          label_col=out_label, medoid=flag, space=space)\n",
    "       \n",
    "    \n",
    "        results[centre+'_matches_species1'] = c_nn1.cross_species_label_matches\n",
    "        results[centre+'_matches_species2'] = c_nn2.cross_species_label_matches\n",
    "        results[centre+'_matches_union'] = len(set(c_nn1.cross_species_label_matches_names).union(\n",
    "                                            set(c_nn2.cross_species_label_matches_names)))\n",
    "    \n",
    "    if ret_matches == True:\n",
    "        return (results, matches, alignments)\n",
    "    \n",
    "    else:\n",
    "        return results\n",
    "    \n",
    "    \n",
    "def get_louvain_metrics(fname, label='cell_type'):\n",
    "    # Compute adjusted rand index for measuring label alignment across species using ground truth information\n",
    "    \n",
    "    adata = sc.read_h5ad(fname)\n",
    "    \n",
    "    try:\n",
    "        sc.pp.pca(adata, n_comps=50)\n",
    "        sc.pp.neighbors(adata, n_neighbors=15)\n",
    "    except:\n",
    "        pass\n",
    "    \n",
    "    if '_' not in adata.obs[label].values[0]:\n",
    "        adata.obs[label].str.cat(adata.obs[\"species\"], sep=\"_\")\n",
    "    \n",
    "    metrics = {}\n",
    "    for resolution in [10, 5, 2, 1, 0.8, 0.5, 0.4, 0.2, 0.1, 0.01, 0.001, 0.0001]:\n",
    "        print('Calculating for resolution: ', str(resolution))\n",
    "        sc.tl.louvain(adata, resolution=resolution)\n",
    "\n",
    "        true_clusters = pd.read_csv('/dfs/project/cross-species/data/lung/shared/true_cell_type_clusters.csv', index_col=0)\n",
    "        true_clusters = true_clusters.merge(adata.obs, left_on='cell_type', right_on=label)\n",
    "        \n",
    "        metrics[resolution] = {\n",
    "            #'ARI':adjusted_rand_score(true_clusters['cluster'], true_clusters['louvain'].astype('int')),\n",
    "            'RI':rand_score(true_clusters['cluster'], true_clusters['louvain'].astype('int')),\n",
    "            'AMI': adjusted_mutual_info_score(true_clusters['cluster'], true_clusters['louvain'].astype('int'))\n",
    "        }\n",
    "\n",
    "    return metrics\n",
    "\n",
    "\n",
    "# Maria's cell type reannotation function\n",
    "def reannotate(adata, source='human', target='mouse', label='cell_type'):\n",
    "    for resolution in [2, 1, 0.8, 0.6, 0.4, 0.2, 0.1]:\n",
    "        sc.tl.louvain(adata, resolution)\n",
    "        louvain_clusters = set(adata.obs['louvain'])\n",
    "\n",
    "        reannotated = {}\n",
    "        for c in louvain_clusters:\n",
    "            current_cluster = adata[adata.obs['louvain']==c]\n",
    "            if len(set(current_cluster.obs['species']))==2:\n",
    "                cluster_source = current_cluster[current_cluster.obs['species']==source]\n",
    "                cluster_target = current_cluster[current_cluster.obs['species']==target]\n",
    "                c = Counter(cluster_source.obs[label])\n",
    "                major_cell_type = max(c.items(), key=operator.itemgetter(1))[0]      \n",
    "                for c in cluster_target.obs_names:\n",
    "                    if c not in reannotated:\n",
    "                        reannotated[c] = major_cell_type\n",
    "    adata_source = adata[adata.obs['species']==source]\n",
    "    tmp = dict(zip(adata_source.obs_names, adata_source.obs[label]))\n",
    "    reannotated = {**reannotated, **tmp}\n",
    "    adata.obs['reannotated_'+source] = [reannotated[c] if c in reannotated else 'None' \n",
    "                                  for c in adata.obs_names]\n",
    "    \n",
    "def get_reannotation_metrics(fname, label='cell_type', source='human', target='mouse'):\n",
    "    # This is current specific to mouse reannotation\n",
    "    \n",
    "    adata = sc.read_h5ad(fname)\n",
    "    sc.pp.neighbors(adata)\n",
    "    \n",
    "    if '_' in adata.obs[label].values[0]:\n",
    "        label = 'labels2'\n",
    "    \n",
    "    reannotate(adata, source=source, label=label)\n",
    "    m = adata[adata.obs['species']==target]\n",
    "    \n",
    "    true_labels_path = '/dfs/project/cross-species/data/lung/shared/true_cell_type.csv'\n",
    "    true_labels = pd.read_csv(true_labels_path, index_col=0)\n",
    "    results_df = true_labels.merge(m.obs, left_on=source+'_cell_type', right_on='reannotated_'+source)\n",
    "    return np.mean(results_df[target+'_cell_type'] == results_df[label])\n",
    "\n",
    "## ---------------------------------\n",
    "## General helper functions\n",
    "## -----------------------------------\n",
    "\n",
    "def plotly_scatter(adata, embed = 'X_umap', label= 'cell_type', \n",
    "                   hover_cols = ['cell_type', 'species']):\n",
    "    plot_df = pd.DataFrame(adata.obsm[embed])\n",
    "    plot_df[label] = adata.obs[label].values\n",
    "    for c in hover_cols:\n",
    "        plot_df[c] = adata.obs[c].values\n",
    "    \n",
    "    fig = px.scatter(plot_df, x=0, y=1, \n",
    "                     hover_name=label,\n",
    "                     color = label,\n",
    "                     hover_data=hover_cols)\n",
    "\n",
    "    fig.show()\n",
    "    \n",
    "def Mode(arr):\n",
    "    # Wrapper for mode\n",
    "    return mode(arr)[0][0]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f992587",
   "metadata": {},
   "source": [
    "# Figure 1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b16181e",
   "metadata": {},
   "source": [
    "### 1A Random Matrices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aab8b2bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(2)\n",
    "y = np.log(np.random.rand(3, 8) ** 1.1 + 1)\n",
    "ax, fig = plt.subplots(figsize=(8, 3))\n",
    "ax = sns.heatmap(y, cmap=\"Blues\", cbar=False);\n",
    "ax.get_xaxis().set_visible(False)\n",
    "ax.get_yaxis().set_visible(False)\n",
    "plt.savefig(\"figures/1a_mat1.svg\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb8710ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(4)\n",
    "y = np.log(np.random.rand(4, 9) ** 1.5 + 1)\n",
    "ax, fig = plt.subplots(figsize=(9, 4))\n",
    "ax = sns.heatmap(y, cmap=\"Oranges\", cbar=False);\n",
    "ax.get_xaxis().set_visible(False)\n",
    "ax.get_yaxis().set_visible(False)\n",
    "plt.savefig(\"figures/1a_mat2.svg\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b26d11f",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(14)\n",
    "y = np.log(np.random.rand(2, 7) ** 1.75 +1)\n",
    "ax, fig = plt.subplots(figsize=(7, 2))\n",
    "ax = sns.heatmap(y, cmap=\"Greens\", cbar=False);\n",
    "ax.get_xaxis().set_visible(False)\n",
    "ax.get_yaxis().set_visible(False)\n",
    "plt.savefig(\"figures/1a_mat3.svg\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1093dcae",
   "metadata": {},
   "source": [
    "# Figure 2 Performance"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d758da9",
   "metadata": {},
   "source": [
    "# Frog / Zebrafish Embryogenesis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2a9b21d",
   "metadata": {},
   "outputs": [],
   "source": [
    "centroids_path = '/dfs/project/cross-species/yanay/data/centroids_seeds/metric_results/test256_data_frog_zebrafish_org_lasso_fz_run_l1_0.0_pe_1.0_ESM2_seed_5.h5ad'\n",
    "samap_path = \"/dfs/project/cross-species/yanay/data/final_embeds/samap_fz_noct.h5ad\"\n",
    "samap_labels_path = '/dfs/project/cross-species/yanay/data/final_embeds/samap_fz_ct.h5ad'\n",
    "scvi_path = \"/dfs/project/cross-species/yanay/data/final_embeds/scvi_fz.h5ad\"\n",
    "scanorama_path = \"/dfs/project/cross-species/yanay/data/scanorama/fz_seeds/seed_0.h5ad\"\n",
    "harmony_path = \"/dfs/project/cross-species/yanay/data/harmony/fz_seeds/seed_0.h5ad\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b01ca7aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "centroids_ad = sc.read(centroids_path)\n",
    "samap_ad = sc.read(samap_path)\n",
    "scvi_ad = sc.read(scvi_path)\n",
    "scanorama_ad = sc.read(scanorama_path)\n",
    "harmony_ad = sc.read(harmony_path)\n",
    "samap_labels_ad = sc.read(samap_labels_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ff1db15",
   "metadata": {},
   "outputs": [],
   "source": [
    "display(centroids_ad, samap_ad, scvi_ad, scanorama_ad, harmony_ad)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a871fe0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Process the UMAPs for each dataset.\n",
    "# We need to do UMAP since that is what SAMAP does only, so need to be fair"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9270c33b",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if \"X_umap\" not in centroids_ad.obsm:\n",
    "    sc.pp.neighbors(centroids_ad, use_rep=\"X\")\n",
    "    sc.tl.umap(centroids_ad, random_state=0)\n",
    "    centroids_ad.write(centroids_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef7a067f",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if \"X_umap\" not in scvi_ad.obsm:\n",
    "    sc.pp.neighbors(scvi_ad, use_rep=\"X\")\n",
    "    sc.tl.umap(scvi_ad, random_state=0)\n",
    "    scvi_ad.write(scvi_path)    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4edad0ad",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if \"X_umap\" not in scanorama_ad.obsm:\n",
    "    sc.pp.neighbors(scanorama_ad, use_rep=\"X\")\n",
    "    sc.tl.umap(scanorama_ad, random_state=0)\n",
    "    scanorama_ad.write(scanorama_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0dfde42f",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if \"X_umap\" not in harmony_ad.obsm:\n",
    "    sc.pp.neighbors(harmony_ad, use_rep=\"X\")\n",
    "    sc.tl.umap(harmony_ad, random_state=0)\n",
    "    harmony_ad.write(harmony_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f54c9a8e",
   "metadata": {},
   "outputs": [],
   "source": [
    "samap_ad.obsm[\"X_umap\"] = samap_ad.X\n",
    "samap_labels_ad.obsm[\"X_umap\"] = samap_labels_ad.X"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e974d00",
   "metadata": {},
   "source": [
    "## 2.A Performance barchart"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c30d877b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# fz paths\n",
    "fz_models_to_paths = {\"SATURN\\n(Our Model)\":\"/dfs/project/cross-species/yanay/data/scoring_csvs/no_lasso_rank_esm2_scores.csv\",#\"/dfs/project/cross-species/yanay/data/scoring_csvs/lasso_fz_scores_scores.csv\", #\"/dfs/project/cross-species/yanay/data/scoring_csvs/fz_30_seeds_scores.csv\",\n",
    "                      #\"Our Model (ESM2)\":\"/dfs/project/cross-species/yanay/data/scoring_csvs/fz_esm2_scores.csv\",\n",
    "                      #\"Our Model (ProtXL)\":\"/dfs/project/cross-species/yanay/data/scoring_csvs/fz_protxl_scores.csv\",\n",
    "                      \"SAMap \\n(Weakly supervised)\":\"/dfs/project/cross-species/yanay/data/scoring_csvs/fz_samap_scores_ct_scores.csv\",\n",
    "                      #\"SAMap\":\"/dfs/project/cross-species/yanay/data/scoring_csvs/fz_samap_noct_scores_scores.csv\",\n",
    "                      \"Harmony\":\"/dfs/project/cross-species/yanay/data/scoring_csvs/harmony_fz_scores.csv\",\n",
    "                      \"scVI\":\"/dfs/project/cross-species/yanay/data/scoring_csvs/fz_scvi_scores_scores.csv\",\n",
    "                      \"Scanorama\":\"/dfs/project/cross-species/yanay/data/scoring_csvs/scanorama_fz_scores.csv\",\n",
    "                      #\"Random\":\"fz_rand_scores.csv\"\n",
    "                     }\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2da97e83",
   "metadata": {},
   "source": [
    "### Figure 2A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1065ba59",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "width = 8 / 6 * len(fz_models_to_paths)\n",
    "fig, ax = plt.subplots(sharex=False, sharey=False, figsize=(width, 3), frameon=False, dpi=300)\n",
    "ax.axhline(y=0.93,  linewidth=2, color='r')\n",
    "ys = np.array([])\n",
    "xs = np.array([])\n",
    "hues = np.array([])\n",
    "\n",
    "\n",
    "for model, path in fz_models_to_paths.items():\n",
    "    if path is not \"\":\n",
    "        scores = pd.read_csv(path)[\"Logistic Regression\"].values\n",
    "        ys = np.append(ys, scores)\n",
    "        xs = np.append(xs, np.array([model] * len(scores)))\n",
    "        hues = np.append(hues, pd.read_csv(path)[\"Label\"].values)\n",
    "        \n",
    "        \n",
    "sns.boxplot(y=ys, x=xs, ax=ax, hue=hues);\n",
    "ax.set_xticklabels(ax.get_xticklabels(), rotation=0, fontsize=10);\n",
    "ax.set_ylim(0, 1);\n",
    "ax.tick_params(axis='both', which='minor', labelsize=7);\n",
    "ax.set_title(\"Model Transfer Accuracy (Frog/Zebrafish)\");\n",
    "\n",
    "title = \"Zebrafish to Frog\"\n",
    "df = pd.DataFrame(columns=[\"x\",\"y\",\"hue\"])\n",
    "df[\"x\"] = xs\n",
    "df[\"y\"] = ys\n",
    "df[\"hue\"] = hues\n",
    "df_1 = df[df[\"hue\"] == title.lower()].drop(columns=\"hue\")\n",
    "M1 = df_1.groupby([\"x\"]).agg(\"median\")\n",
    "SE1 = df_1.groupby([\"x\"]).agg(np.std, ddof=1)\n",
    "\n",
    "fig, ax = plt.subplots(sharex=False, sharey=False, figsize=(width, 3), frameon=False, dpi=300)\n",
    "ax.axhline(y=0.93,  linewidth=2, color='r', linestyle=\"--\")\n",
    "\n",
    "ax.set_xticks(ticks=np.arange(len(M1.index)), labels=M1.index)\n",
    "ax.set_xticklabels(fz_models_to_paths.keys(), rotation=0, fontsize=10);\n",
    "ax.grid(False)\n",
    "\n",
    "for x in fz_models_to_paths.keys():\n",
    "    ax = plt.errorbar(x=x, y=M1.loc[x].values.flatten(), yerr=(SE1.loc[x].values).flatten(), fmt='o', ls=\"none\", label=x)\n",
    "\n",
    "fig.legend().remove()\n",
    "plt.ylim(0,1)\n",
    "plt.title(\"Model Label Transfer Accuracy \" + f\"({title})\")\n",
    "plt.savefig(\"figures/2a.svg\")\n",
    "plt.show()\n",
    "\n",
    "\n",
    "title = \"Frog to Zebrafish\"\n",
    "df_2 = df[df[\"hue\"] == title.lower()].drop(columns=\"hue\")\n",
    "M2 = df_2.groupby([\"x\"]).agg(\"mean\")\n",
    "SE2 = df_2.groupby([\"x\"]).agg(np.std, ddof=1)\n",
    "\n",
    "\n",
    "fig, ax = plt.subplots(sharex=False, sharey=False, figsize=(width, 3), frameon=False, dpi=300)\n",
    "ax.axhline(y=0.93,  linewidth=2, color='r', linestyle=\"--\")\n",
    "\n",
    "ax.set_xticks(ticks=np.arange(len(M1.index)), labels=M1.index)\n",
    "ax.set_xticklabels(fz_models_to_paths.keys(), rotation=0, fontsize=10);\n",
    "ax.grid(False)\n",
    "for x in fz_models_to_paths.keys():\n",
    "    ax = plt.errorbar(x=x, y=M2.loc[x].values.flatten(), yerr=(SE2.loc[x].values).flatten(), fmt='o', ls=\"none\", label=x)\n",
    "\n",
    "fig.legend().remove()\n",
    "plt.ylim(0,1)\n",
    "plt.title(\"Model Label Transfer Accuracy \" + f\"({title})\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c921c5ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "M1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c3d61fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "M2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b523c62",
   "metadata": {},
   "outputs": [],
   "source": [
    "b = M1.loc[\"SAMap \\n(Weakly supervised)\"]\n",
    "a = M1.loc[\"SATURN\\n(Our Model)\"]\n",
    "print(f\"{int((a-b)/b * 100)}% improvement\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43fa6a31",
   "metadata": {},
   "source": [
    "## Figure 2B Stacked Bars"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "318c547e",
   "metadata": {},
   "outputs": [],
   "source": [
    "cen_al_f = alignment_score(centroids_path,  col=\"labels2\", species=\"frog\")\n",
    "sam_al_f = alignment_score(samap_labels_path,  col=\"cell_type\", species=\"frog\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce937915",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Most common cell types\n",
    "labels_by_rank = pd.DataFrame(centroids_ad.obs[\"labels2\"].value_counts()).reset_index()\n",
    "labels_by_rank"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b7a8f89",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_cluster_knn_bar(df, source='human', other='mouse', ax=None, title=None):\n",
    "    \"\"\"\n",
    "    Creates a stacked bar plot to identify majority k nearest neighbors for\n",
    "    a given cluster, on a specific axis\n",
    "    \"\"\"\n",
    "    if ax is None:\n",
    "        fig, (ax) = plt.subplots(1, 1, sharex=True, sharey=True, figsize=(6, 17), frameon=False)\n",
    "    bars = defaultdict(int)\n",
    "    colors = defaultdict(int)\n",
    "    tick = -1; \n",
    "    tick_pos = {}\n",
    "    colors_list = ['','green'] + [\"#DEDEDE\"]*100\n",
    "    df = df.sort_values(['labels2', 'Correct', 'Score'], ascending=[True, False, False])\n",
    "    \n",
    "    for i in df.iterrows():\n",
    "        color = 'k'\n",
    "        x = i[1]['Source_Cluster']\n",
    "        y = i[1]['Score']        \n",
    "\n",
    "        if i[1]['Correct']:\n",
    "            color = 'green'\n",
    "        else:\n",
    "            color = \"white\"\n",
    "        left = bars[x]\n",
    "        bars[x] = bars[x] + y\n",
    "        colors[x] = colors[x] + 1\n",
    "        if colors_list[colors[x]] == 'green':    \n",
    "            tick = tick+1\n",
    "            tick_pos[x] = tick\n",
    "        ax.barh(tick_pos[x], y, left=left, color=color, alpha=0.5, edgecolor=\"none\")\n",
    "        \n",
    "    keys = sorted(tick_pos.keys())\n",
    "    vals = [tick_pos[k] for k in keys]\n",
    "    ax.set_yticks(vals, keys)\n",
    "    ax.set_ylabel(source)\n",
    "    #ax.set_xlabel('Percentage of cells with cross-species KNN class')\n",
    "    if title is not None:\n",
    "        ax.set_title(title)\n",
    "    ax.set_xlim(0, 1)\n",
    "    return "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b7a7fcb",
   "metadata": {},
   "outputs": [],
   "source": [
    "ct_map = pd.read_csv(\"/dfs/project/cross-species/yanay/fz_true_ct.csv\").set_index(\"Unnamed: 0\").reset_index(drop=True)\n",
    "ct_map.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c5adc7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Most common cell types\n",
    "labels_by_rank = pd.DataFrame(centroids_ad.obs[\"labels2\"].value_counts()).reset_index()\n",
    "labels_by_rank_zf = pd.DataFrame(centroids_ad[centroids_ad.obs[\"species\"] == \"zebrafish\"].obs[\"labels2\"].value_counts()).reset_index()\n",
    "labels_by_rank = labels_by_rank.merge(labels_by_rank_zf.drop(columns=\"labels2\"), left_on=\"index\", right_on=\"index\", how=\"inner\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "031820bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "labels_by_rank"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d33ce6f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "cen_al_f, sam_al_f, cen_al_z, sam_al_z = c1.copy(), c2.copy(), c3.copy(), c4.copy()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7793a345",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "width = 9\n",
    "fig, (row1, row2) = plt.subplots(2, 2, sharex=True, sharey=False, figsize=(width, 5), frameon=False)\n",
    "# Most Common\n",
    "ax1, ax2 = row1\n",
    "ax3, ax4 = row2\n",
    "\n",
    "sam_most = sam_al_f.merge(ct_map, left_on=\"Source_Cluster\", right_on=\"frog_cell_type\")\n",
    "cen_most = cen_al_f.merge(ct_map, left_on=\"Source_Cluster\", right_on=\"frog_cell_type\")\n",
    "\n",
    "sam_most[\"Correct\"] = sam_most[\"Cross_Species_KNN_Label\"] == sam_most[\"zebrafish_cell_type\"]\n",
    "cen_most[\"Correct\"] = cen_most[\"Cross_Species_KNN_Label\"] == cen_most[\"zebrafish_cell_type\"]\n",
    "\n",
    "\n",
    "cen_most = cen_most.merge(labels_by_rank, left_on=\"frog_cell_type\", right_on=\"index\")\n",
    "sam_most = sam_most.merge(labels_by_rank, left_on=\"frog_cell_type\", right_on=\"index\")\n",
    "\n",
    "cen_most = cen_most[cen_most[\"labels2\"] > 10000].sort_values([\"labels2\", \"Source_Cluster\"], ascending=False)\n",
    " # top 5\n",
    "sam_most = sam_most[sam_most[\"labels2\"] > 10000].sort_values([\"labels2\", \"Source_Cluster\"], ascending=False)\n",
    " # top 5\n",
    "\n",
    "\n",
    "\n",
    "plot_cluster_knn_bar(sam_most, source=\"frog\", ax=ax2, title=\"SAMap (Weakly supervised)\")\n",
    "plot_cluster_knn_bar(cen_most, source=\"frog\", ax=ax1, title=\"SATURN (Our Model)\")\n",
    "\n",
    "ax1.set(ylabel=None)\n",
    "ax1.grid(False)\n",
    "ax2.grid(False)\n",
    "ax2.get_yaxis().set_visible(False)\n",
    "ax2.get_xaxis().set_visible(False)\n",
    "ax1.get_xaxis().set_visible(False)\n",
    "\n",
    "# Least Common\n",
    "sam_least = sam_al_f.merge(ct_map, left_on=\"Source_Cluster\", right_on=\"frog_cell_type\")\n",
    "cen_least = cen_al_f.merge(ct_map, left_on=\"Source_Cluster\", right_on=\"frog_cell_type\")\n",
    "\n",
    "\n",
    "sam_least[\"Correct\"] = sam_least[\"Cross_Species_KNN_Label\"] == sam_least[\"zebrafish_cell_type\"]\n",
    "cen_least[\"Correct\"] = cen_least[\"Cross_Species_KNN_Label\"] == cen_least[\"zebrafish_cell_type\"]\n",
    "\n",
    "least_common_modifiers = { \n",
    "                          \"Olfactory placode\":[\"Olfactory placode\"],\n",
    "                          \"Germline\":[\"Germline\"],\n",
    "                          \"Hatching gland\":[\"Hatching gland\"]\n",
    "                         }\n",
    "sam_least_correct = []\n",
    "for row in sam_least.iterrows():\n",
    "    row = row[1]\n",
    "    pred = row['Cross_Species_KNN_Label']\n",
    "    src = row[\"Source_Cluster\"]\n",
    "    res = row[\"zebrafish_cell_type\"]\n",
    "    \n",
    "    sam_least_correct.append((pred == res) or pred in least_common_modifiers.get(src, []))\n",
    "    #1/0\n",
    "                             \n",
    "sam_least[\"Correct\"] = sam_least_correct\n",
    "\n",
    "cen_least_correct = []\n",
    "for row in cen_least.iterrows():\n",
    "    row = row[1]\n",
    "    pred = row['Cross_Species_KNN_Label']\n",
    "    src = row[\"Source_Cluster\"]\n",
    "    res = row[\"zebrafish_cell_type\"]\n",
    "    \n",
    "    cen_least_correct.append((pred == res) or pred in least_common_modifiers.get(src, []))\n",
    "cen_least[\"Correct\"] = cen_least_correct\n",
    "\n",
    "\n",
    "cen_least = cen_least.merge(labels_by_rank, left_on=\"frog_cell_type\", right_on=\"index\")\n",
    "sam_least = sam_least.merge(labels_by_rank, left_on=\"frog_cell_type\", right_on=\"index\")\n",
    "\n",
    "cen_least = cen_least[cen_least[\"labels2\"] < 500].sort_values([\"labels2\", \"Source_Cluster\"], ascending=False)\n",
    " # top 5\n",
    "sam_least = sam_least[sam_least[\"labels2\"] < 500].sort_values([\"labels2\", \"Source_Cluster\"], ascending=False)\n",
    " # top 5\n",
    "\n",
    "plot_cluster_knn_bar(sam_least, source=\"frog\", ax=ax4, title=\"\")\n",
    "plot_cluster_knn_bar(cen_least, source=\"frog\", ax=ax3, title=\"\")\n",
    "\n",
    "ax3.set(ylabel=None)\n",
    "ax3.grid(False)\n",
    "ax4.grid(False)\n",
    "ax4.get_yaxis().set_visible(False);\n",
    "\n",
    "ax3.set(title=None);\n",
    "ax4.set(title=None);\n",
    "\n",
    "\n",
    "fig.tight_layout()\n",
    "\n",
    "plt.savefig(\"figures/2b.svg\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36f1cbe1",
   "metadata": {},
   "source": [
    "### Check Germline Cells"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "196280e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "cen_al_f[cen_al_f[\"Source_Cluster\"] == \"Germline\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69c8ef10",
   "metadata": {},
   "outputs": [],
   "source": [
    "sam_al_f[sam_al_f[\"Source_Cluster\"] == \"Germline\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a551701",
   "metadata": {},
   "source": [
    "## Figure 2D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "396e4307",
   "metadata": {},
   "outputs": [],
   "source": [
    "sc._settings.settings._vector_friendly=True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc7a5d49",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "fig, (axB, axA) = plt.subplots(2, 5, sharex=False, sharey=False, figsize=(13.3, 4), frameon=False, dpi=300)\n",
    "\n",
    "J = 0\n",
    "# Our method\n",
    "\n",
    "# SATURN\n",
    "ax1 = axA[J]\n",
    "fig.tight_layout()\n",
    "sc.pl.umap(centroids_ad, color=\"species\", title=\"\", ax=ax1, show=False)\n",
    "ax1.get_legend().remove();\n",
    "ax1.set(xlabel=None, ylabel=None);\n",
    "ax1.title.set_size(18);\n",
    "\n",
    "ax2 = axB[J]\n",
    "sc.pl.umap(centroids_ad, color=\"labels2\", title=\"SATURN\", ax=ax2, show=False)\n",
    "ax2.get_legend().remove();\n",
    "ax2.set(xlabel=None, ylabel=None);\n",
    "ax2.title.set_size(18);\n",
    "fig.tight_layout()\n",
    "\n",
    "J += 1\n",
    "\n",
    "\n",
    "# SAMAP Label Share\n",
    "ax1 = axA[J]\n",
    "fig.tight_layout()\n",
    "sc.pl.umap(samap_labels_ad, color=\"species\", title=\"\", ax=ax1, show=False)\n",
    "ax1.get_legend().remove();\n",
    "ax1.set(xlabel=None, ylabel=None);\n",
    "ax1.title.set_size(18);\n",
    "\n",
    "ax2 = axB[J]\n",
    "sc.pl.umap(samap_labels_ad, color=\"cell_type\", title=\"SAMap\\n(Weakly supervised)\", ax=ax2, show=False)\n",
    "ax2.get_legend().remove();\n",
    "ax2.set(xlabel=None, ylabel=None);\n",
    "ax2.title.set_size(18);\n",
    "fig.tight_layout()\n",
    "\n",
    "J += 1\n",
    "\n",
    "# Harmony\n",
    "ax1 = axA[J]\n",
    "fig.tight_layout()\n",
    "sc.pl.umap(harmony_ad, color=\"species\", title=\"\", ax=ax1, show=False)\n",
    "ax1.get_legend().remove();\n",
    "ax1.set(xlabel=None, ylabel=None);\n",
    "ax1.title.set_size(18);\n",
    "\n",
    "ax2 = axB[J]\n",
    "sc.pl.umap(harmony_ad, color=\"cell_type\", title=\"Harmony\", ax=ax2, show=False)\n",
    "ax2.get_legend().remove();\n",
    "ax2.set(xlabel=None, ylabel=None);\n",
    "ax2.title.set_size(18);\n",
    "fig.tight_layout()\n",
    "\n",
    "J += 1\n",
    "\n",
    "# SCVI\n",
    "ax1 = axA[J]\n",
    "fig.tight_layout()\n",
    "sc.pl.umap(scvi_ad, color=\"species\", title=\"\", ax=ax1, show=False)\n",
    "ax1.get_legend().remove();\n",
    "ax1.set(xlabel=None, ylabel=None);\n",
    "ax1.title.set_size(18);\n",
    "\n",
    "ax2 = axB[J]\n",
    "sc.pl.umap(scvi_ad, color=\"cell_type\", title=\"scVI\", ax=ax2, show=False)\n",
    "ax2.get_legend().remove();\n",
    "ax2.set(xlabel=None, ylabel=None);\n",
    "ax2.title.set_size(18);\n",
    "fig.tight_layout()\n",
    "\n",
    "J += 1\n",
    "\n",
    "\n",
    "# Scanorama\n",
    "ax1 = axA[J]\n",
    "fig.tight_layout()\n",
    "sc.pl.umap(scanorama_ad, color=\"species\", title=\"\", ax=ax1, show=False)\n",
    "ax1.get_legend().remove();\n",
    "ax1.set(xlabel=None, ylabel=None);\n",
    "ax1.title.set_size(18);\n",
    "\n",
    "ax2 = axB[J]\n",
    "sc.pl.umap(scanorama_ad, color=\"cell_type\", title=\"Scanorama\", ax=ax2, show=False)\n",
    "ax2.get_legend().remove();\n",
    "ax2.set(xlabel=None, ylabel=None);\n",
    "ax2.title.set_size(18);\n",
    "fig.tight_layout()\n",
    "\n",
    "J += 1\n",
    "\n",
    "\n",
    "plt.savefig(\"figures/2d.svg\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a6f624f",
   "metadata": {},
   "outputs": [],
   "source": [
    "samap_ad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7cf50c3",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "plotly_scatter(samap_labels_ad, label=\"cell_type\", hover_cols=[\"cell_type\", \"species\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e613f5f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "plotly_scatter(centroids_ad, label=\"labels2\", hover_cols=[\"labels2\", \"species\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e594682",
   "metadata": {},
   "outputs": [],
   "source": [
    "(centroids_ad.obs[\"labels\"].str.contains(\"frog_Germline\")).sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15ab2d88",
   "metadata": {},
   "outputs": [],
   "source": [
    "(centroids_ad.obs[\"labels\"].str.contains(\"zebrafish_Germline\")).sum()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66a50962",
   "metadata": {},
   "source": [
    "# 2C FZ Outsets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4435c7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.set_figure_params(figsize=(4, 4))\n",
    "\n",
    "# for white background of figures (only for docs rendering)\n",
    "%config InlineBackend.print_figure_kwargs={'facecolor' : \"w\"}\n",
    "%config InlineBackend.figure_format='retina'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ebdac8f",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "centroids_ad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95f7fccd",
   "metadata": {},
   "outputs": [],
   "source": [
    "top10 = list(centroids_ad.obs[\"labels2\"].value_counts()[:10].index)\n",
    "\n",
    "labels_hidden_dict = {a:\"_\" +a for  a in centroids_ad.obs[\"labels2\"].unique()}\n",
    "for l in top10:\n",
    "    labels_hidden_dict[l] = l\n",
    "\n",
    "mapped_labels = [labels_hidden_dict[l] for l in centroids_ad.obs[\"labels2\"]]\n",
    "centroids_ad.obs[\"hidden_labels\"] = mapped_labels\n",
    "\n",
    "spaced_labels = [l.replace(\" \", \"\\n\").title() for l in centroids_ad.obs[\"labels2\"]]\n",
    "centroids_ad.obs[\"spaced_labels\"] = spaced_labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c30dfe40",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(sharex=False, sharey=False, figsize=(4, 5), frameon=False, dpi=300)\n",
    "\n",
    "to_show = [a.replace(' ', '\\n').title() for a in ['Blastula', '', 'Neuroectoderm', 'Non-neural ectoderm', '', \n",
    "           '', '', 'Epidermal progenitor', \n",
    "           '', '', '', '', \n",
    "           'Tailbud', 'Intermediate mesoderm', '', 'Eye primordium', \n",
    "           'Placodal area', 'Hindbrain', 'Neural crest', 'Neuron', 'Forebrain/midbrain', \n",
    "           '', '', '', '', \n",
    "           'Skeletal muscle', '', 'Blood', '', '', \n",
    "           '', '', '', '', \n",
    "           '', 'Optic', 'Pluripotent', '', '', \n",
    "           '', '', '', '', \n",
    "           '', '', '', \n",
    "           'Macrophage', 'Spemann Organizer', 'Goblet cell', 'Presomitic mesoderm']]\n",
    "ax = sc.pl.umap(centroids_ad, color=\"spaced_labels\", show=False, alpha=1, legend_loc='on data', ax=ax, legend_fontsize=7, legend_fontoutline=2)\n",
    "ax.set(xlabel=None, ylabel=None);\n",
    "ax.set(title=None)\n",
    "for a in ax.texts:\n",
    "    if not (a._text in to_show):\n",
    "        #print(a._text)\n",
    "        a.set_visible(False)\n",
    "    if a._text == 'Spemann\\nOrganizer':\n",
    "        a.set_text(\"1\")\n",
    "    if a._text == 'Goblet\\nCell':\n",
    "        a.set_text(\"2\")\n",
    "    if a._text == 'Macrophage':\n",
    "        a.set_text(\"3\")\n",
    "ax.spines['top'].set_visible(False)\n",
    "ax.spines['right'].set_visible(False)\n",
    "ax.spines['bottom'].set_visible(False)\n",
    "ax.spines['left'].set_visible(False)\n",
    "plt.savefig(\"figures/2c_main.svg\")    \n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8c17ca2",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(sharex=False, sharey=False, figsize=(4, 5), frameon=False, dpi=300)\n",
    "\n",
    "ax = sc.pl.umap(centroids_ad, color=\"species\", show=False, alpha=1, ax=ax, legend_loc=None)\n",
    "ax.set(xlabel=None, ylabel=None);\n",
    "ax.set(title=None)\n",
    "ax.spines['top'].set_visible(False)\n",
    "ax.spines['right'].set_visible(False)\n",
    "ax.spines['bottom'].set_visible(False)\n",
    "ax.spines['left'].set_visible(False)\n",
    "plt.savefig(\"figures/2c_spec.svg\")\n",
    "plt.show()\n",
    "xlims = ax.get_xlim()\n",
    "ylims = ax.get_ylim()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ae7f4bd",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(sharex=False, sharey=False, figsize=(4, 4), frameon=False, dpi=300)\n",
    "\n",
    "ax = sc.pl.umap(centroids_ad, color=\"species\", show=False, alpha=1, ax=ax, legend_loc=None, groups=\"zebrafish\")\n",
    "ax.set(xlabel=None, ylabel=None);\n",
    "ax.set(title=None)\n",
    "ax.spines['top'].set_visible(False)\n",
    "ax.spines['right'].set_visible(False)\n",
    "ax.spines['bottom'].set_visible(False)\n",
    "ax.spines['left'].set_visible(False)\n",
    "plt.savefig(\"figures/2c_zebrafish.svg\")\n",
    "plt.show()\n",
    "xlims = ax.get_xlim()\n",
    "ylims = ax.get_ylim()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "462c9e1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "mod = 1.4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be5238d1",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(sharex=False, sharey=False, figsize=(4/mod, 4/mod), frameon=False, dpi=300)\n",
    "\n",
    "zoom = 6\n",
    "x_offset = -10\n",
    "y_offset = 9\n",
    "xlim = (np.array(xlims) - min(xlims))/zoom + x_offset\n",
    "ylim = (np.array(ylims) - min(ylims))/zoom + y_offset\n",
    "\n",
    "\n",
    "umap_limited = centroids_ad[(centroids_ad.obsm[\"X_umap\"][:, 0] > xlim[0]) & (centroids_ad.obsm[\"X_umap\"][:, 0] < xlim[1]) \\\n",
    "& (centroids_ad.obsm[\"X_umap\"][:, 1] > ylim[0]) & (centroids_ad.obsm[\"X_umap\"][:, 1] > ylim[0])]\n",
    "ax = sc.pl.umap(umap_limited, color=\"labels2\", show=False, title=\"Inset 1\", alpha=1, ax=ax, legend_fontsize=6, legend_fontoutline=2, groups=[\"Macrophage\", \"Myeloid progenitors\", \"Blood\"])\n",
    "\n",
    "for a in ax.get_legend().texts:\n",
    "    a._text = a._text.title()\n",
    "    if a._text == \"Na\":\n",
    "        a._text = \"Other\"\n",
    "plt.xlim(*xlim);\n",
    "plt.ylim(*ylim);\n",
    "ax.set(title=None, xlabel=None, ylabel=None);\n",
    "plt.savefig(\"figures/2c_3.svg\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8fd83833",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(sharex=False, sharey=False, figsize=(4/mod, 4/mod), frameon=False, dpi=300)\n",
    "\n",
    "zoom = 12\n",
    "x_offset = 5.25\n",
    "y_offset = 0\n",
    "xlim = (np.array(xlims) - min(xlims))/zoom + x_offset\n",
    "ylim = (np.array(ylims) - min(ylims))/zoom + y_offset\n",
    "\n",
    "\n",
    "umap_limited = centroids_ad[(centroids_ad.obsm[\"X_umap\"][:, 0] > xlim[0]) & (centroids_ad.obsm[\"X_umap\"][:, 0] < xlim[1]) \\\n",
    "& (centroids_ad.obsm[\"X_umap\"][:, 1] > ylim[0]) & (centroids_ad.obsm[\"X_umap\"][:, 1] > ylim[0])]\n",
    "ax = sc.pl.umap(umap_limited, color=\"labels2\", show=False, title=\"Inset 1\", alpha=1, ax=ax, legend_fontsize=6, legend_fontoutline=2, groups=[\"Ionocyte\", \"Goblet cell\",\n",
    "                                                                                                                                            \"Periderm\", \"Rare epidermal subtypes\",\n",
    "                                                                                                                                            \"Epidermal progenitor\"])\n",
    "\n",
    "\n",
    "for a in ax.get_legend().texts:\n",
    "    a._text = a._text.title()\n",
    "    if a._text == \"Na\":\n",
    "        a._text = \"Other\"\n",
    "plt.xlim(*xlim);\n",
    "plt.ylim(*ylim);\n",
    "ax.set(title=None, xlabel=None, ylabel=None);\n",
    "plt.savefig(\"figures/2c_2.svg\")\n",
    "\n",
    "#plt.show()\n",
    "#fig, ax = plt.subplots(sharex=False, sharey=False, figsize=(6, 4), frameon=False, dpi=300)\n",
    "\n",
    "\n",
    "\n",
    "#ax = sc.pl.umap(umap_limited, color=\"species\", show=False, title=\"Inset 1\", alpha=1, legend_loc='on data', ax=ax, legend_fontsize=6, legend_fontoutline=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "536f0d61",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(sharex=False, sharey=False, figsize=(4/mod, 4/mod), frameon=False, dpi=300)\n",
    "\n",
    "zoom = 5\n",
    "x_offset = 8.25\n",
    "y_offset = 5.5\n",
    "xlim = (np.array(xlims) - min(xlims))/zoom + x_offset\n",
    "ylim = (np.array(ylims) - min(ylims))/zoom + y_offset\n",
    "\n",
    "\n",
    "umap_limited = centroids_ad[(centroids_ad.obsm[\"X_umap\"][:, 0] > xlim[0]) & (centroids_ad.obsm[\"X_umap\"][:, 0] < xlim[1]) \\\n",
    "& (centroids_ad.obsm[\"X_umap\"][:, 1] > ylim[0]) & (centroids_ad.obsm[\"X_umap\"][:, 1] > ylim[0])]\n",
    "ax = sc.pl.umap(umap_limited, color=\"labels2\", show=False, title=\"Inset 1\", alpha=1, ax=ax, legend_fontsize=6, legend_fontoutline=2, groups=[\"Spemann organizer\", \"Involuting marginal zone\",\n",
    "                                                                                                                                            \"Dorsal organizer\", \n",
    "                                                                                                                                            \"Endoderm\"])\n",
    "\n",
    "\n",
    "for a in ax.get_legend().texts:\n",
    "    a._text = a._text.title()\n",
    "    if a._text == \"Na\":\n",
    "        a._text = \"Other\"\n",
    "plt.xlim(*xlim);\n",
    "plt.ylim(*ylim);\n",
    "ax.set(title=None, xlabel=None, ylabel=None);\n",
    "plt.savefig(\"figures/2c_1.svg\")\n",
    "#plt.show()\n",
    "#fig, ax = plt.subplots(sharex=False, sharey=False, figsize=(6, 4), frameon=False, dpi=300)\n",
    "\n",
    "\n",
    "\n",
    "#ax = sc.pl.umap(umap_limited, color=\"species\", show=False, title=\"Inset 1\", alpha=1, legend_loc='on data', ax=ax, legend_fontsize=6, legend_fontoutline=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6b24296",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(sharex=False, sharey=False, figsize=(4/mod, 4/mod), frameon=False, dpi=300)\n",
    "\n",
    "#zoom = 5\n",
    "#x_offset = 8.25\n",
    "#y_offset = 5.5\n",
    "#xlim = (np.array(xlims) - min(xlims))/zoom + x_offset\n",
    "#ylim = (np.array(ylims) - min(ylims))/zoom + y_offset\n",
    "\n",
    "\n",
    "#map_limited = centroids_ad[(centroids_ad.obsm[\"X_umap\"][:, 0] > xlim[0]) & (centroids_ad.obsm[\"X_umap\"][:, 0] < xlim[1]) \\\n",
    "#& (centroids_ad.obsm[\"X_umap\"][:, 1] > ylim[0]) & (centroids_ad.obsm[\"X_umap\"][:, 1] > ylim[0])]\n",
    "ax = sc.pl.umap(centroids_ad, color=\"species\", show=False, title=\"Inset 4\", alpha=1, ax=ax, legend_fontsize=6, legend_fontoutline=2)\n",
    "#plt.ylim(*ylim);\n",
    "ax.set(title=None, xlabel=None, ylabel=None);\n",
    "plt.savefig(\"figures/2c_4.svg\")\n",
    "#plt.show()\n",
    "#fig, ax = plt.subplots(sharex=False, sharey=False, figsize=(6, 4), frameon=False, dpi=300)\n",
    "\n",
    "\n",
    "\n",
    "#ax = sc.pl.umap(umap_limited, color=\"species\", show=False, title=\"Inset 1\", alpha=1, legend_loc='on data', ax=ax, legend_fontsize=6, legend_fontoutline=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6cacdd69",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "6688aa64",
   "metadata": {},
   "source": [
    "# Biohub Poster Figure 2 Excerpts "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6556188",
   "metadata": {},
   "outputs": [],
   "source": [
    "width = 15 # inches\n",
    "height = 4\n",
    "\n",
    "fig, (axB, axA) = plt.subplots(2, 5, sharex=False, sharey=False, figsize=(width, height), frameon=False, dpi=300)\n",
    "\n",
    "J = 0\n",
    "# Our method\n",
    "\n",
    "# SATURN\n",
    "ax1 = axA[J]\n",
    "fig.tight_layout()\n",
    "sc.pl.umap(centroids_ad, color=\"species\", title=\"\", ax=ax1, show=False)\n",
    "ax1.get_legend().remove();\n",
    "ax1.set(xlabel=None, ylabel=None);\n",
    "ax1.title.set_size(18);\n",
    "\n",
    "ax2 = axB[J]\n",
    "sc.pl.umap(centroids_ad, color=\"labels2\", title=\"SATURN\", ax=ax2, show=False)\n",
    "ax2.get_legend().remove();\n",
    "ax2.set(xlabel=None, ylabel=None);\n",
    "ax2.title.set_size(18);\n",
    "fig.tight_layout()\n",
    "\n",
    "J += 1\n",
    "\n",
    "\n",
    "# SAMAP Label Share\n",
    "ax1 = axA[J]\n",
    "fig.tight_layout()\n",
    "sc.pl.umap(samap_labels_ad, color=\"species\", title=\"\", ax=ax1, show=False)\n",
    "ax1.get_legend().remove();\n",
    "ax1.set(xlabel=None, ylabel=None);\n",
    "ax1.title.set_size(18);\n",
    "\n",
    "ax2 = axB[J]\n",
    "sc.pl.umap(samap_labels_ad, color=\"cell_type\", title=\"SAMap\\n(Weakly supervised)\", ax=ax2, show=False)\n",
    "ax2.get_legend().remove();\n",
    "ax2.set(xlabel=None, ylabel=None);\n",
    "ax2.title.set_size(18);\n",
    "fig.tight_layout()\n",
    "\n",
    "J += 1\n",
    "\n",
    "'''\n",
    "# SAMAP No Label Share\n",
    "ax1 = axA[J]\n",
    "fig.tight_layout()\n",
    "sc.pl.umap(samap_ad, color=\"species\", title=\"\", ax=ax1, show=False)\n",
    "ax1.get_legend().remove();\n",
    "ax1.set(xlabel=None, ylabel=None);\n",
    "ax1.title.set_size(18);\n",
    "\n",
    "ax2 = axB[J]\n",
    "sc.pl.umap(samap_ad, color=\"cell_type\", title=\"SAMap\", ax=ax2, show=False)\n",
    "ax2.get_legend().remove();\n",
    "ax2.set(xlabel=None, ylabel=None);\n",
    "ax2.title.set_size(18);\n",
    "fig.tight_layout()\n",
    "\n",
    "J += 1\n",
    "'''\n",
    "\n",
    "# Harmony\n",
    "ax1 = axA[J]\n",
    "fig.tight_layout()\n",
    "sc.pl.umap(harmony_ad, color=\"species\", title=\"\", ax=ax1, show=False)\n",
    "ax1.get_legend().remove();\n",
    "ax1.set(xlabel=None, ylabel=None);\n",
    "ax1.title.set_size(18);\n",
    "\n",
    "ax2 = axB[J]\n",
    "sc.pl.umap(harmony_ad, color=\"cell_type\", title=\"Harmony\", ax=ax2, show=False)\n",
    "ax2.get_legend().remove();\n",
    "ax2.set(xlabel=None, ylabel=None);\n",
    "ax2.title.set_size(18);\n",
    "fig.tight_layout()\n",
    "\n",
    "J += 1\n",
    "\n",
    "# SCVI\n",
    "ax1 = axA[J]\n",
    "fig.tight_layout()\n",
    "sc.pl.umap(scvi_ad, color=\"species\", title=\"\", ax=ax1, show=False)\n",
    "ax1.get_legend().remove();\n",
    "ax1.set(xlabel=None, ylabel=None);\n",
    "ax1.title.set_size(18);\n",
    "\n",
    "ax2 = axB[J]\n",
    "sc.pl.umap(scvi_ad, color=\"cell_type\", title=\"scVI\", ax=ax2, show=False)\n",
    "ax2.get_legend().remove();\n",
    "ax2.set(xlabel=None, ylabel=None);\n",
    "ax2.title.set_size(18);\n",
    "fig.tight_layout()\n",
    "\n",
    "J += 1\n",
    "\n",
    "\n",
    "# Scanorama\n",
    "ax1 = axA[J]\n",
    "fig.tight_layout()\n",
    "sc.pl.umap(scanorama_ad, color=\"species\", title=\"\", ax=ax1, show=False)\n",
    "ax1.get_legend().remove();\n",
    "ax1.set(xlabel=None, ylabel=None);\n",
    "ax1.title.set_size(18);\n",
    "\n",
    "ax2 = axB[J]\n",
    "sc.pl.umap(scanorama_ad, color=\"cell_type\", title=\"Scanorama\", ax=ax2, show=False)\n",
    "ax2.get_legend().remove();\n",
    "ax2.set(xlabel=None, ylabel=None);\n",
    "ax2.title.set_size(18);\n",
    "fig.tight_layout()\n",
    "\n",
    "J += 1\n",
    "\n",
    "\n",
    "plt.savefig(\"figures/2d_poster.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e810cd90",
   "metadata": {},
   "outputs": [],
   "source": [
    "width = 7.5\n",
    "height = 7.5\n",
    "fig, ax = plt.subplots(sharex=False, sharey=False, figsize=(width, height), frameon=False, dpi=300)\n",
    "\n",
    "to_show = [a.replace(' ', '\\n').title() for a in ['Blastula', '', 'Neuroectoderm', 'Non-neural ectoderm', '', \n",
    "           '', '', 'Epidermal progenitor', \n",
    "           '', '', '', '', \n",
    "           'Tailbud', 'Intermediate mesoderm', '', 'Eye primordium', \n",
    "           'Placodal area', 'Hindbrain', 'Neural crest', 'Neuron', 'Forebrain/midbrain', \n",
    "           '', '', '', '', \n",
    "           'Skeletal muscle', '', 'Blood', '', '', \n",
    "           '', '', '', '', \n",
    "           '', 'Optic', 'Pluripotent', '', '', \n",
    "           '', '', '', '', \n",
    "           '', '', '', \n",
    "           'Macrophage', 'Spemann Organizer', 'Goblet cell', 'Presomitic mesoderm']]\n",
    "ax = sc.pl.umap(centroids_ad, color=\"spaced_labels\", show=False, alpha=1, legend_loc='on data', ax=ax, legend_fontsize=7, legend_fontoutline=2)\n",
    "ax.set(xlabel=None, ylabel=None);\n",
    "ax.set(title=None)\n",
    "for a in ax.texts:\n",
    "    if not (a._text in to_show):\n",
    "        #print(a._text)\n",
    "        a.set_visible(False)\n",
    "    if a._text == 'Spemann\\nOrganizer':\n",
    "        a.set_text(\"\")\n",
    "    if a._text == 'Goblet\\nCell':\n",
    "        a.set_text(\"\")\n",
    "    if a._text == 'Macrophage':\n",
    "        a.set_text(\"\")\n",
    "    a.set_fontsize(12)\n",
    "ax.spines['top'].set_visible(False)\n",
    "ax.spines['right'].set_visible(False)\n",
    "ax.spines['bottom'].set_visible(False)\n",
    "ax.spines['left'].set_visible(False)\n",
    "plt.savefig(\"figures/2c_main_poster.svg\")    \n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa7cb1bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "width = 7.5 / 1.75\n",
    "height = 4 / 1.75\n",
    "fig, ax = plt.subplots(sharex=False, sharey=False, figsize=(width, height), frameon=False, dpi=300)\n",
    "ax.axhline(y=0.93,  linewidth=2, color='r')\n",
    "ys = np.array([])\n",
    "xs = np.array([])\n",
    "hues = np.array([])\n",
    "\n",
    "\n",
    "for model, path in fz_models_to_paths.items():\n",
    "    if path is not \"\":\n",
    "        scores = pd.read_csv(path)[\"Logistic Regression\"].values\n",
    "        ys = np.append(ys, scores)\n",
    "        xs = np.append(xs, np.array([model] * len(scores)))\n",
    "        hues = np.append(hues, pd.read_csv(path)[\"Label\"].values)\n",
    "        \n",
    "        \n",
    "sns.boxplot(y=ys, x=xs, ax=ax, hue=hues);\n",
    "ax.set_xticklabels(ax.get_xticklabels(), rotation=0, fontsize=10);\n",
    "ax.set_ylim(0, 1);\n",
    "ax.tick_params(axis='both', which='minor', labelsize=7);\n",
    "ax.set_title(\"Model Transfer Accuracy (Frog/Zebrafish)\");\n",
    "\n",
    "title = \"Zebrafish to Frog\"\n",
    "df = pd.DataFrame(columns=[\"x\",\"y\",\"hue\"])\n",
    "df[\"x\"] = xs\n",
    "df[\"y\"] = ys\n",
    "df[\"hue\"] = hues\n",
    "df_1 = df[df[\"hue\"] == title.lower()].drop(columns=\"hue\")\n",
    "M1 = df_1.groupby([\"x\"]).agg(\"median\")\n",
    "SE1 = df_1.groupby([\"x\"]).agg(np.std, ddof=1)\n",
    "\n",
    "fig, ax = plt.subplots(sharex=False, sharey=False, figsize=(width, height), frameon=False, dpi=300)\n",
    "ax.axhline(y=0.93,  linewidth=2, color='r', linestyle=\"--\")\n",
    "\n",
    "ax.set_xticks(ticks=np.arange(len(M1.index)), labels=M1.index)\n",
    "ax.set_xticklabels([x.split(\"\\n\")[0] for x in fz_models_to_paths.keys()], rotation=0, fontsize=10);\n",
    "ax.grid(False)\n",
    "\n",
    "for x in fz_models_to_paths.keys():\n",
    "    ax = plt.errorbar(x=x, y=M1.loc[x].values.flatten(), yerr=(SE1.loc[x].values).flatten(), fmt='o', ls=\"none\", label=x.split(\"\\n\")[0], lw=4, ms=8)\n",
    "\n",
    "fig.legend().remove()\n",
    "plt.ylim(0,1)\n",
    "#plt.title(\"Label Transfer Accuracy \" + f\"({title})\")\n",
    "plt.savefig(\"figures/2a_poster.png\")\n",
    "plt.show()\n",
    "\n",
    "\n",
    "title = \"Frog to Zebrafish\"\n",
    "df_2 = df[df[\"hue\"] == title.lower()].drop(columns=\"hue\")\n",
    "M2 = df_2.groupby([\"x\"]).agg(\"mean\")\n",
    "SE2 = df_2.groupby([\"x\"]).agg(np.std, ddof=1)\n",
    "\n",
    "\n",
    "fig, ax = plt.subplots(sharex=False, sharey=False, figsize=(width, height), frameon=False, dpi=300)\n",
    "ax.axhline(y=0.93,  linewidth=2, color='r', linestyle=\"--\")\n",
    "\n",
    "ax.set_xticks(ticks=np.arange(len(M1.index)), labels=M1.index)\n",
    "ax.set_xticklabels([x.split(\"\\n\")[0] for x in fz_models_to_paths.keys()], rotation=0, fontsize=10);\n",
    "ax.grid(False)\n",
    "for x in fz_models_to_paths.keys():\n",
    "    ax = plt.errorbar(x=x, y=M2.loc[x].values.flatten(), yerr=(SE2.loc[x].values).flatten(), fmt='o', ls=\"none\", label=x.split(\"\\n\")[0], lw=3)\n",
    "\n",
    "fig.legend().remove()\n",
    "plt.ylim(0,1)\n",
    "#plt.title(\"Label Transfer Accuracy \" + f\"({title})\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5f2a7ef",
   "metadata": {},
   "source": [
    "# Figure 3"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0e02b95",
   "metadata": {},
   "source": [
    "## Figure 3B\n",
    "\n",
    "Example subset 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fbb2287b",
   "metadata": {},
   "outputs": [],
   "source": [
    "subset_1_types = [\"Macrophage\", \"Myeloid progenitors\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97b378a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "subset_1_ad = centroids_ad[centroids_ad.obs[\"labels2\"].isin(subset_1_types)]\n",
    "sc.pl.umap(subset_1_ad, color=\"labels2\")\n",
    "sc.pl.umap(subset_1_ad, color=\"species\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "238b4887",
   "metadata": {},
   "outputs": [],
   "source": [
    "ancs_ad = sc.AnnData(centroids_ad.obsm[\"ancs\"])\n",
    "ancs_ad.obs = centroids_ad.obs\n",
    "ancs_ad.obs[\"subset\"] = centroids_ad.obs[\"labels2\"].isin(subset_1_types).astype(str)\n",
    "\n",
    "ancs_ad.obsm[\"X_umap\"] = centroids_ad.obsm[\"X_umap\"]\n",
    "ancs_ad.X = np.log(ancs_ad.X+1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0fbad09f",
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.tl.dendrogram(centroids_ad,groupby='labels2', use_rep='X')\n",
    "ancs_ad.uns[\"dendrogram_labels2\"] = centroids_ad.uns[\"dendrogram_labels2\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5d81c39",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(centroids_path.replace(\".h5ad\", \"_genes_to_centroids.pkl\"),'rb') as f:\n",
    "    genes_to_scores = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10763ffe",
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.tl.rank_genes_groups(ancs_ad, 'subset', groups=[\"True\"], method='wilcoxon')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72840f54",
   "metadata": {},
   "outputs": [],
   "source": [
    "best_gene_cluster_1 = sc.get.rank_genes_groups_df(ancs_ad, group=\"True\")[\"names\"][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3aaa037",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_scores(anc_idx):\n",
    "    '''\n",
    "    Given the index of a centroid, return the scores by gene for that centroid\n",
    "    '''\n",
    "    scores = {}\n",
    "    for (gene), embs in genes_to_scores.items():\n",
    "        scores[gene] = embs[anc_idx]\n",
    "    return scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42dbfc0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "best_gene_cluster_1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f5e266e",
   "metadata": {},
   "outputs": [],
   "source": [
    "de_df = sc.get.rank_genes_groups_df(ancs_ad, group=[\"True\"])\n",
    "de_df_pos = de_df[de_df[\"scores\"] > 0]\n",
    "de_df_neg = de_df[de_df[\"scores\"] < 0].sort_values(\"scores\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2abe30e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "de_df_pos.head(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed755e2a",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "for de in de_df_pos.head(6)[\"names\"]:\n",
    "    max_gene = pd.DataFrame(get_scores(int(de)).items(), columns=[\"gene\", \"weight\"]).sort_values(\"weight\", ascending=False)\n",
    "    print(de)\n",
    "    display(max_gene.head(20))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4383f04a",
   "metadata": {},
   "outputs": [],
   "source": [
    "subset_1_groups = [\"Macrophage\", \"Myeloid progenitors\", \"Blood\", \"Periderm\", \n",
    "                   \"Endothelial\", \"Pluripotent\"]\n",
    "\n",
    "subset_1_ad = ancs_ad[ancs_ad.obs[\"labels2\"].isin(subset_1_groups)]\n",
    "subset_1_ad_embs = centroids_ad[centroids_ad.obs[\"labels2\"].isin(subset_1_groups)]\n",
    "\n",
    "sc.tl.dendrogram(subset_1_ad_embs, groupby='labels2', use_rep='X')\n",
    "subset_1_ad.uns[\"dendrogram_labels2\"] = subset_1_ad_embs.uns[\"dendrogram_labels2\"]\n",
    "markers = [str(a) for a in de_df_pos.head(6)[\"names\"] if a != \"1302\"]\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(6,3), dpi=300)\n",
    "dp = sc.pl.dotplot(subset_1_ad, markers, groupby='labels2', \n",
    "                   dendrogram=True, show=False, cmap=\"viridis\",\n",
    "                   ax=ax, vmin=0.0, vmax=0.5, dot_min=0, dot_max=1,\n",
    "                  )\n",
    "\n",
    "dp[\"mainplot_ax\"].set_xticklabels([\"Arhgdi\", \"Cebp\", \"Ptp\", \"Cybb\", \"Lcp\"], rotation = 0, ha=\"center\", fontsize=8)\n",
    "dp[\"mainplot_ax\"].set_yticklabels(dp[\"mainplot_ax\"].get_yticklabels(), fontsize=8)\n",
    "plt.savefig(\"figures/3b_dot.svg\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "650c467f",
   "metadata": {},
   "source": [
    "# Figure 3C\n",
    "### Subset 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8852cab",
   "metadata": {},
   "outputs": [],
   "source": [
    "subset_2_types = [\"Ionocyte\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e61d4466",
   "metadata": {},
   "outputs": [],
   "source": [
    "subset_2_ad = centroids_ad[centroids_ad.obs[\"labels2\"].isin(subset_2_types)]\n",
    "sc.pl.umap(subset_2_ad, color=\"labels2\")\n",
    "sc.pl.umap(subset_2_ad, color=\"species\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ef4d946",
   "metadata": {},
   "outputs": [],
   "source": [
    "subset_2_ad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42fdec98",
   "metadata": {},
   "outputs": [],
   "source": [
    "ancs_ad = sc.AnnData(centroids_ad.obsm[\"ancs\"])\n",
    "ancs_ad.obs = centroids_ad.obs\n",
    "ancs_ad.obs[\"subset\"] = centroids_ad.obs[\"labels2\"].isin(subset_2_types).astype(str)\n",
    "\n",
    "ancs_ad.obsm[\"X_umap\"] = centroids_ad.obsm[\"X_umap\"]\n",
    "ancs_ad.X = np.log(ancs_ad.X+1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93ac019e",
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.tl.dendrogram(centroids_ad, groupby='labels2', use_rep='X')\n",
    "ancs_ad.uns[\"dendrogram_labels2\"] = centroids_ad.uns[\"dendrogram_labels2\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0f22460",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(centroids_path.replace(\".h5ad\", \"_genes_to_centroids.pkl\"),'rb') as f:\n",
    "    genes_to_scores = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b44c7260",
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.tl.rank_genes_groups(ancs_ad, 'subset', groups=[\"True\"], method='wilcoxon')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b1912f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "best_gene_cluster_2 = sc.get.rank_genes_groups_df(ancs_ad, group=\"True\")[\"names\"][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c232313e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_scores(anc_idx):\n",
    "    '''\n",
    "    Given the index of a centroid, return the scores by gene for that centroid\n",
    "    '''\n",
    "    scores = {}\n",
    "    for (gene), embs in genes_to_scores.items():\n",
    "        scores[gene] = embs[anc_idx]\n",
    "    return scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e476cea",
   "metadata": {},
   "outputs": [],
   "source": [
    "best_gene_cluster_2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "945bc247",
   "metadata": {},
   "outputs": [],
   "source": [
    "de_df = sc.get.rank_genes_groups_df(ancs_ad, group=[\"True\"])\n",
    "de_df_pos = de_df[de_df[\"scores\"] > 0]\n",
    "de_df_neg = de_df[de_df[\"scores\"] < 0].sort_values(\"scores\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eedbff20",
   "metadata": {},
   "outputs": [],
   "source": [
    "de_df_pos.head(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "487e682e",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "for de in de_df_pos.head(6)[\"names\"]:\n",
    "    print(de)\n",
    "    max_gene = pd.DataFrame(get_scores(int(de)).items(), columns=[\"gene\", \"weight\"]).sort_values(\"weight\", ascending=False)\n",
    "    display(max_gene.head(10))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08317c8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "ionocyte_groups = [\"Ionocyte\", \"Small secretory cells\", \"Rare epidermal subtypes\", \"Goblet cell\", \n",
    "                   \"Endothelial\", \"Secretory epidermal\"]\n",
    "\n",
    "subset_2_ad = ancs_ad[ancs_ad.obs[\"labels2\"].isin(ionocyte_groups)]\n",
    "subset_2_ad_embs = centroids_ad[centroids_ad.obs[\"labels2\"].isin(ionocyte_groups)]\n",
    "\n",
    "sc.tl.dendrogram(subset_2_ad_embs, groupby='labels2', use_rep='X')\n",
    "subset_2_ad.uns[\"dendrogram_labels2\"] = subset_2_ad_embs.uns[\"dendrogram_labels2\"]\n",
    "markers = [str(a) for a in de_df_pos.head(6)[\"names\"] if a != \"1302\"]\n",
    "fig, ax = plt.subplots(figsize=(6,3), dpi=300)\n",
    "dp = sc.pl.dotplot(subset_2_ad, markers, groupby='labels2', dendrogram=True, show=False, cmap=\"viridis\", ax=ax, vmin=0.0, vmax=0.5, dot_min=0, dot_max=1)\n",
    "\n",
    "dp[\"mainplot_ax\"].set_xticklabels([\"Foxi\", \"Dmr2\", \"Cldn\", \"Ubp\", \"Atp6v0\"], rotation = 0, ha=\"center\", fontsize=8)\n",
    "dp[\"mainplot_ax\"].set_yticklabels(dp[\"mainplot_ax\"].get_yticklabels(), fontsize=8)\n",
    "plt.savefig(\"figures/3c_dot.svg\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "828be6f6",
   "metadata": {},
   "source": [
    "# Figure 3D"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19f841f2",
   "metadata": {},
   "source": [
    "## now compare these CTs between species"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c968826",
   "metadata": {},
   "outputs": [],
   "source": [
    "subset_2_ad = centroids_ad[centroids_ad.obs[\"labels2\"].isin(subset_2_types)]\n",
    "sc.pp.pca(subset_2_ad)\n",
    "sc.pp.neighbors(subset_2_ad)\n",
    "sc.tl.umap(subset_2_ad)\n",
    "sc.pl.umap(subset_2_ad, color=\"labels2\")\n",
    "sc.pl.umap(subset_2_ad, color=\"species\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6057a2ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "subset_2_ancs_ad = ancs_ad[ancs_ad.obs[\"labels2\"].isin(subset_2_types)]\n",
    "subset_2_ancs_ad.obsm[\"X_umap\"] = subset_2_ad.obsm[\"X_umap\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8c1b54f",
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.tl.rank_genes_groups(subset_2_ancs_ad, 'species', method='wilcoxon')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4a522d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "best_gene_cluster = sc.get.rank_genes_groups_df(subset_2_ancs_ad, group=\"zebrafish\")[\"names\"][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "899898b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_scores(anc_idx):\n",
    "    '''\n",
    "    Given the index of a centroid, return the scores by gene for that centroid\n",
    "    '''\n",
    "    scores = {}\n",
    "    for (gene), embs in genes_to_scores.items():\n",
    "        scores[gene] = embs[anc_idx]\n",
    "    return scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "589479b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "de_df = sc.get.rank_genes_groups_df(subset_2_ancs_ad, group=[\"zebrafish\"])\n",
    "de_df_pos = de_df[de_df[\"scores\"] > 0]\n",
    "de_df_neg = de_df[de_df[\"scores\"] < 0].sort_values(\"scores\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48250a4a",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "for de in de_df_pos.head(5)[\"names\"]:\n",
    "    max_gene = pd.DataFrame(get_scores(int(de)).items(), columns=[\"gene\", \"weight\"]).sort_values(\"weight\", ascending=False)\n",
    "    print(de)\n",
    "    display(max_gene.head(7))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29ccd05e",
   "metadata": {},
   "outputs": [],
   "source": [
    "de_df = sc.get.rank_genes_groups_df(subset_2_ancs_ad, group=[\"frog\"])\n",
    "de_df_neg = de_df[de_df[\"scores\"] > 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "671378aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "markers = list(de_df_pos.head(20)[\"names\"].astype(str)) + list(de_df_neg.head(20)[\"names\"].astype(str))\n",
    "sc.pl.dotplot(ancs_ad, markers, groupby='species', dendrogram=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d63ddb91",
   "metadata": {},
   "outputs": [],
   "source": [
    "#fair_markers = [\"1090\", \"463\", \"157\", \"1160\", \"1627\", \"1064\"]\n",
    "#fair_markers = [\"1372\", \"13\", \"1861\"]\n",
    "fair_markers = [\"613\", \"501\", \"148\"]\n",
    "\n",
    "print(\"All Cell Types\")\n",
    "sc.pl.dotplot(ancs_ad, fair_markers, groupby='species', dendrogram=True, title=\"All Cell Types\")\n",
    "print(\"Ionocytes Only\")\n",
    "\n",
    "sc.pl.dotplot(subset_2_ancs_ad, fair_markers, groupby='species', dendrogram=True, title=\"Ionocytes Only\")\n",
    "sc.pl.umap(subset_2_ancs_ad, color=fair_markers)\n",
    "sc.pl.umap(subset_2_ancs_ad, color=\"species\", title=\"Ionocytes Only\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6db58d36",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "for de in fair_markers:\n",
    "    max_gene = pd.DataFrame(get_scores(int(de)).items(), columns=[\"gene\", \"weight\"]).sort_values(\"weight\", ascending=False)\n",
    "    print(de)\n",
    "    display(max_gene.head(7))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d274a7e1",
   "metadata": {},
   "source": [
    "## Make Figure 3d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98ff6ad8",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(5,3.5))\n",
    "\n",
    "#fair_markers = [\"1090\", \"463\", \"157\", \"1160\", \"1627\", \"1064\"]\n",
    "#fair_markers = [\"1372\", \"13\", \"1861\"]\n",
    "\n",
    "fair_markers = [\"613\", \"501\", \"148\"]\n",
    "i=0\n",
    "sc.pl.umap(subset_2_ancs_ad, color=f\"{fair_markers[i]}\", title=f\"Macrogene: {fair_markers[i]}\", ax=ax1, show=False, colorbar_loc=None, vmin=0);\n",
    "i+=1\n",
    "sc.pl.umap(subset_2_ancs_ad, color=f\"{fair_markers[i]}\", title=f\"Macrogene: {fair_markers[i]}\", ax=ax2, show=False, colorbar_loc=None, vmin=0);\n",
    "i+=1\n",
    "sc.pl.umap(subset_2_ancs_ad, color=f\"{fair_markers[i]}\", title=f\"Macrogene: {fair_markers[i]}\", ax=ax3, show=False, colorbar_loc=None, vmin=0);\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "sc.pl.umap(subset_2_ancs_ad, color=\"species\", ax=ax4, show=False, title=\"Species\", legend_loc=None);\n",
    "\n",
    "ax1.set(xlabel=None, ylabel=None);\n",
    "ax2.set(xlabel=None, ylabel=None);\n",
    "ax3.set(xlabel=None, ylabel=None);\n",
    "ax4.set(xlabel=None, ylabel=None);\n",
    "\n",
    "plt.savefig(\"figures/3d.svg\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "046b4080",
   "metadata": {},
   "outputs": [],
   "source": [
    "for de in fair_markers:\n",
    "    max_gene = pd.DataFrame(get_scores(int(de)).items(), columns=[\"gene\", \"weight\"]).sort_values(\"weight\", ascending=False)\n",
    "    print(de)\n",
    "    display(max_gene.head(7))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c7a1d58",
   "metadata": {},
   "source": [
    "# Supplement Figures"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47aef683",
   "metadata": {},
   "source": [
    "# Supplement Figures 1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b5cd9f3",
   "metadata": {},
   "source": [
    "## Supplement 1A Bar Graphs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68c7b17f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_cluster_knn_bar(df, source='human', other='mouse', ax=None, title=None):\n",
    "    \"\"\"\n",
    "    Creates a stacked bar plot to identify majority k nearest neighbors for\n",
    "    a given cluster, on a specific axis\n",
    "    \"\"\"\n",
    "    if ax is None:\n",
    "        fig, (ax) = plt.subplots(1, 1, sharex=True, sharey=True, figsize=(6, 17), frameon=False)\n",
    "    bars = defaultdict(int)\n",
    "    colors = defaultdict(int)\n",
    "    tick = -1; \n",
    "    tick_pos = {}\n",
    "    colors_list = ['','green'] + [\"#DEDEDE\"]*100\n",
    "    df = df.sort_values(['labels2', 'Correct', 'Score'], ascending=[True, False, False])\n",
    "    \n",
    "    for i in df.iterrows():\n",
    "        color = 'k'\n",
    "        x = i[1]['Source_Cluster']\n",
    "        y = i[1]['Score']        \n",
    "\n",
    "        if i[1]['Correct']:\n",
    "            color = 'green'\n",
    "        else:\n",
    "            color = \"white\"\n",
    "        left = bars[x]\n",
    "        bars[x] = bars[x] + y\n",
    "        colors[x] = colors[x] + 1\n",
    "        if colors_list[colors[x]] == 'green':    \n",
    "            tick = tick+1\n",
    "            tick_pos[x] = tick\n",
    "        ax.barh(tick_pos[x], y, left=left, color=color, alpha=0.5, edgecolor=\"none\")\n",
    "        \n",
    "    keys = sorted(tick_pos.keys())\n",
    "    vals = [tick_pos[k] for k in keys]\n",
    "    ax.set_yticks(vals, keys)\n",
    "    ax.set_ylabel(source)\n",
    "    #ax.set_xlabel('Percentage of cells with cross-species KNN class')\n",
    "    if title is not None:\n",
    "        ax.set_title(title)\n",
    "    ax.set_xlim(0, 1)\n",
    "    return "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9dbd907",
   "metadata": {},
   "outputs": [],
   "source": [
    "cen_al_f = alignment_score(centroids_path,  col=\"labels2\", species=\"frog\")\n",
    "sam_al_f = alignment_score(samap_labels_path,  col=\"cell_type\", species=\"frog\")\n",
    "\n",
    "cen_al_z = alignment_score(centroids_path,  col=\"labels2\", species=\"zebrafish\")\n",
    "sam_al_z = alignment_score(samap_labels_path,  col=\"cell_type\", species=\"zebrafish\")\n",
    "c1, c2, c3, c4 = cen_al_f.copy(), sam_al_f.copy(), cen_al_z.copy(), sam_al_z.copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2830ed6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "ct_map = pd.read_csv(\"/dfs/project/cross-species/yanay/fz_true_ct.csv\").set_index(\"Unnamed: 0\").reset_index(drop=True)\n",
    "ct_map.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f422eaf2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Most common cell types\n",
    "labels_by_rank_fg = pd.DataFrame(centroids_ad[centroids_ad.obs[\"species\"] == \"frog\"].obs[\"labels2\"].value_counts()).reset_index()\n",
    "labels_by_rank_zf = pd.DataFrame(centroids_ad[centroids_ad.obs[\"species\"] == \"zebrafish\"].obs[\"labels2\"].value_counts()).reset_index()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce8f95b5",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "cen_al_f, sam_al_f, cen_al_z, sam_al_z = c1.copy(), c2.copy(), c3.copy(), c4.copy()\n",
    "width = 9\n",
    "height = 16\n",
    "#fig, (ax1, ax2, ax3, ax4) = plt.subplots(1, 4, sharex=True, sharey=False, figsize=(width, height), frameon=False)\n",
    "fig, (row1, row2) = plt.subplots(2, 2, sharex=True, sharey=False, figsize=(width, height), frameon=False)\n",
    "# Most Common\n",
    "ax1, ax2 = row1\n",
    "ax3, ax4 = row2\n",
    "\n",
    "# Frog\n",
    "\n",
    "sam_al_f = sam_al_f.merge(ct_map, left_on=\"Source_Cluster\", right_on=\"frog_cell_type\")\n",
    "cen_al_f = cen_al_f.merge(ct_map, left_on=\"Source_Cluster\", right_on=\"frog_cell_type\")\n",
    "\n",
    "\n",
    "sam_al_f[\"Correct\"] = sam_al_f[\"Cross_Species_KNN_Label\"] == sam_al_f[\"zebrafish_cell_type\"]\n",
    "cen_al_f[\"Correct\"] = cen_al_f[\"Cross_Species_KNN_Label\"] == cen_al_f[\"zebrafish_cell_type\"]\n",
    "\n",
    "least_common_modifiers = {\n",
    "                          \"Olfactory placode\":[\"Olfactory placode\"],\n",
    "                          \"Germline\":[\"Germline\"],\n",
    "                          \"Hatching gland\":[\"Hatching gland\"]\n",
    "                         }\n",
    "sam_al_f_correct = []\n",
    "for row in sam_al_f.iterrows():\n",
    "    row = row[1]\n",
    "    pred = row['Cross_Species_KNN_Label']\n",
    "    src = row[\"Source_Cluster\"]\n",
    "    res = row[\"zebrafish_cell_type\"]\n",
    "    \n",
    "    sam_al_f_correct.append((pred == res) or pred in least_common_modifiers.get(src, []))\n",
    "    #1/0\n",
    "                             \n",
    "sam_al_f[\"Correct\"] = sam_al_f_correct\n",
    "\n",
    "cen_al_f_correct = []\n",
    "for row in cen_al_f.iterrows():\n",
    "    row = row[1]\n",
    "    pred = row['Cross_Species_KNN_Label']\n",
    "    src = row[\"Source_Cluster\"]\n",
    "    res = row[\"zebrafish_cell_type\"]\n",
    "    \n",
    "    cen_al_f_correct.append((pred == res) or pred in least_common_modifiers.get(src, []))\n",
    "cen_al_f[\"Correct\"] = cen_al_f_correct\n",
    "\n",
    "\n",
    "cen_al_f = cen_al_f.merge(labels_by_rank_zf, left_on=\"frog_cell_type\", right_on=\"index\")\n",
    "sam_al_f = sam_al_f.merge(labels_by_rank_zf, left_on=\"frog_cell_type\", right_on=\"index\")\n",
    "\n",
    "cen_al_f = cen_al_f.sort_values([\"labels2\", \"Source_Cluster\"], ascending=False)\n",
    " # top 5\n",
    "sam_al_f = sam_al_f.sort_values([\"labels2\", \"Source_Cluster\"], ascending=False)\n",
    " # top 5\n",
    "\n",
    "plot_cluster_knn_bar(sam_al_f, source=\"frog\", ax=ax2, title=\"\")\n",
    "plot_cluster_knn_bar(cen_al_f, source=\"frog\", ax=ax1, title=\"\")\n",
    "\n",
    "ax1.set(ylabel=None)\n",
    "ax1.grid(False)\n",
    "ax2.grid(False)\n",
    "ax2.get_yaxis().set_visible(False);\n",
    "\n",
    "ax1.set(title=None);\n",
    "ax2.set(title=None);\n",
    "\n",
    "# Zebrafish\n",
    "sam_al_z = sam_al_z.merge(ct_map, left_on=\"Source_Cluster\", right_on=\"zebrafish_cell_type\")\n",
    "cen_al_z = cen_al_z.merge(ct_map, left_on=\"Source_Cluster\", right_on=\"zebrafish_cell_type\")\n",
    "\n",
    "\n",
    "sam_al_z[\"Correct\"] = sam_al_z[\"Cross_Species_KNN_Label\"] == sam_al_z[\"frog_cell_type\"]\n",
    "cen_al_z[\"Correct\"] = cen_al_z[\"Cross_Species_KNN_Label\"] == cen_al_z[\"frog_cell_type\"]\n",
    "\n",
    "least_common_modifiers = {\n",
    "                          \"Olfactory placode\":[\"Olfactory placode\"],\n",
    "                          \"Germline\":[\"Germline\"],\n",
    "                          \"Hatching gland\":[\"Hatching gland\"]\n",
    "                         }\n",
    "sam_al_z_correct = []\n",
    "for row in sam_al_z.iterrows():\n",
    "    row = row[1]\n",
    "    pred = row['Cross_Species_KNN_Label']\n",
    "    src = row[\"Source_Cluster\"]\n",
    "    res = row[\"frog_cell_type\"]\n",
    "    \n",
    "    sam_al_z_correct.append((pred == res) or pred in least_common_modifiers.get(src, []))\n",
    "    #1/0\n",
    "                             \n",
    "sam_al_z[\"Correct\"] = sam_al_z_correct\n",
    "\n",
    "cen_al_z_correct = []\n",
    "for row in cen_al_z.iterrows():\n",
    "    row = row[1]\n",
    "    pred = row['Cross_Species_KNN_Label']\n",
    "    src = row[\"Source_Cluster\"]\n",
    "    res = row[\"frog_cell_type\"]\n",
    "    \n",
    "    cen_al_z_correct.append((pred == res) or pred in least_common_modifiers.get(src, []))\n",
    "cen_al_z[\"Correct\"] = cen_al_z_correct\n",
    "\n",
    "\n",
    "cen_al_z = cen_al_z.merge(labels_by_rank_fg, left_on=\"zebrafish_cell_type\", right_on=\"index\")\n",
    "sam_al_z = sam_al_z.merge(labels_by_rank_fg, left_on=\"zebrafish_cell_type\", right_on=\"index\")\n",
    "\n",
    "cen_al_z = cen_al_z.sort_values([\"labels2\", \"Source_Cluster\"], ascending=False)\n",
    " # top 5\n",
    "sam_al_z = sam_al_z.sort_values([\"labels2\", \"Source_Cluster\"], ascending=False)\n",
    " # top 5\n",
    "\n",
    "plot_cluster_knn_bar(sam_al_z, source=\"zebrafish\", ax=ax4, title=\"\")\n",
    "plot_cluster_knn_bar(cen_al_z, source=\"zebrafish\", ax=ax3, title=\"\")\n",
    "\n",
    "ax3.set(ylabel=None)\n",
    "ax3.grid(False)\n",
    "ax4.grid(False)\n",
    "ax4.get_yaxis().set_visible(False);\n",
    "\n",
    "ax3.set(title=None);\n",
    "ax4.set(title=None);\n",
    "\n",
    "\n",
    "fig.tight_layout()\n",
    "\n",
    "plt.savefig(\"figures/supp1_bar.svg\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7d4c69c",
   "metadata": {},
   "source": [
    "## Supplement 1B Bar Chart"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6f7d065",
   "metadata": {},
   "outputs": [],
   "source": [
    "fz_models_to_paths_sup = {\n",
    "#\"Both Losses\\n(ESM1b)\":\"/dfs/project/cross-species/yanay/data/scoring_csvs/lasso_fz_scores_scores.csv\",\n",
    "#\"L1, no Rank Loss\\n(ESM1b)\":\"/dfs/project/cross-species/yanay/data/scoring_csvs/lasso_no_rank_scores.csv\",\n",
    "\n",
    "#\"(ESM1b)\":\"/dfs/project/cross-species/yanay/data/scoring_csvs/fz_30_seeds_scores.csv\",\n",
    "\"ESM2\":\"/dfs/project/cross-species/yanay/data/scoring_csvs/no_lasso_rank_esm2_scores.csv\",\n",
    "\"ESM1b\":\"/dfs/project/cross-species/yanay/data/scoring_csvs/no_lasso_rank_scores.csv\",\n",
    "#\"ESM2\":\"/dfs/project/cross-species/yanay/data/scoring_csvs/fz_esm2_scores.csv\",\n",
    "\"ProtXL\":\"/dfs/project/cross-species/yanay/data/scoring_csvs/protxl_new_loss_scores.csv\",\n",
    "\"ESM2 Pretrain\":\"/dfs/project/cross-species/yanay/data/scoring_csvs/no_lasso_rank_esm2_pretrain_scores.csv\",\n",
    "\"SAMAP \\n(Weakly Supervised)\":\"/dfs/project/cross-species/yanay/data/scoring_csvs/fz_samap_scores_ct_scores.csv\",\n",
    "\"SAMAP \\n(Unsupervised)\":\"/dfs/project/cross-species/yanay/data/scoring_csvs/fz_samap_noct_scores_scores.csv\",\n",
    "\"scVI\":\"/dfs/project/cross-species/yanay/data/scoring_csvs/fz_scvi_scores_scores.csv\",\n",
    "    \n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "771f2a8c",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "width = len(fz_models_to_paths_sup) * 1.25\n",
    "fig, ax = plt.subplots(sharex=False, sharey=False, figsize=(width, 3), frameon=False, dpi=300)\n",
    "\n",
    "ys = np.array([])\n",
    "xs = np.array([])\n",
    "hues = np.array([])\n",
    "\n",
    "\n",
    "for model, path in fz_models_to_paths_sup.items():\n",
    "    if path is not \"\":\n",
    "        scores = pd.read_csv(path)[\"Logistic Regression\"].values\n",
    "        ys = np.append(ys, scores)\n",
    "        xs = np.append(xs, np.array([model] * len(scores)))\n",
    "        hues = np.append(hues, pd.read_csv(path)[\"Label\"].values)\n",
    "        \n",
    "        print(f\"{path}: Max: {max(scores)}\")\n",
    "sns.boxplot(y=ys, x=xs, ax=ax, hue=hues);\n",
    "ax.set_xticklabels(ax.get_xticklabels(), rotation=0, fontsize=10);\n",
    "#ax.set_ylim(0, 1);\n",
    "ax.tick_params(axis='both', which='minor', labelsize=7);\n",
    "ax.set_title(\"Label Transfer Accuracy\");\n",
    "plt.savefig(\"figures/supp1b_box.svg\")\n",
    "title = \"Zebrafish to Frog\"\n",
    "df = pd.DataFrame(columns=[\"x\",\"y\",\"hue\"])\n",
    "\n",
    "df[\"x\"] = xs\n",
    "df[\"y\"] = ys\n",
    "df[\"hue\"] = hues\n",
    "df_1 = df[df[\"hue\"] == title.lower()].drop(columns=\"hue\")\n",
    "M1 = df_1.groupby([\"x\"]).agg(\"median\")\n",
    "SE1 = df_1.groupby([\"x\"]).agg(np.std, ddof=1)\n",
    "\n",
    "fig, ax = plt.subplots(sharex=False, sharey=False, figsize=(width, 3), frameon=False, dpi=300)\n",
    "ax.set_xticks(ticks=np.arange(len(M1.index)), labels=M1.index)\n",
    "ax.set_xticklabels(fz_models_to_paths_sup.keys(), rotation=0, fontsize=10);\n",
    "ax.grid(False)\n",
    "\n",
    "for x in fz_models_to_paths_sup.keys():\n",
    "    ax = plt.errorbar(x=x, y=M1.loc[x].values.flatten(), yerr=(SE1.loc[x].values).flatten(), fmt='o', ls=\"none\", label=x)\n",
    "\n",
    "fig.legend().remove()\n",
    "#plt.ylim(0,1)\n",
    "plt.title(\"Accuracy of Transferring Labels from Zebrafish to Frog\")\n",
    "\n",
    "plt.show()\n",
    "\n",
    "\n",
    "title = \"Frog to Zebrafish\"\n",
    "df_2 = df[df[\"hue\"] == title.lower()].drop(columns=\"hue\")\n",
    "M2 = df_2.groupby([\"x\"]).agg(\"mean\")\n",
    "SE2 = df_2.groupby([\"x\"]).agg(np.std, ddof=1)\n",
    "\n",
    "\n",
    "fig, ax = plt.subplots(sharex=False, sharey=False, figsize=(width, 3), frameon=False, dpi=300)\n",
    "ax.set_xticks(ticks=np.arange(len(M1.index)), labels=M1.index)\n",
    "ax.set_xticklabels(fz_models_to_paths_sup.keys(), rotation=0, fontsize=10);\n",
    "ax.grid(False)\n",
    "for x in fz_models_to_paths_sup.keys():\n",
    "    ax = plt.errorbar(x=x, y=M2.loc[x].values.flatten(), yerr=(SE2.loc[x].values).flatten(), fmt='o', ls=\"none\", label=x)\n",
    "\n",
    "fig.legend().remove()\n",
    "#plt.ylim(0,1)\n",
    "plt.title(\"Accuracy of Transferring Labels from Frog to Zebrafish\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c514f63",
   "metadata": {},
   "outputs": [],
   "source": [
    "pd.concat([M1, M2, SE1, SE2], axis=1).set_axis([\"Median Z to F\", \"Median F to Z\", \"SE Z to F\", \"SE F to Z\"], axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91582d79",
   "metadata": {},
   "outputs": [],
   "source": [
    "M2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e1d833d",
   "metadata": {},
   "outputs": [],
   "source": [
    "SE1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb92ceae",
   "metadata": {},
   "outputs": [],
   "source": [
    "SE2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e9a4134",
   "metadata": {},
   "outputs": [],
   "source": [
    "display(df_2.groupby([\"x\"]).agg(\"max\").sort_values(\"y\"))\n",
    "display(df_1.groupby([\"x\"]).agg(\"max\").sort_values(\"y\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97da8aab",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
